82 lines
3.5 KiB
Python
82 lines
3.5 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from engine_compat import QWEN_CHAT_TEMPLATE, CompatibilityError, prepare_model
|
|
|
|
|
|
class EngineCompatTests(unittest.TestCase):
|
|
def _model(self, root: Path, tokenizer: dict) -> Path:
|
|
model = root / "model"
|
|
model.mkdir()
|
|
(model / "config.json").write_text(
|
|
json.dumps({"model_type": "qwen2", "architectures": ["Qwen2ForCausalLM"]}),
|
|
encoding="utf-8",
|
|
)
|
|
(model / "tokenizer_config.json").write_text(
|
|
json.dumps(tokenizer), encoding="utf-8"
|
|
)
|
|
(model / "model.safetensors").write_bytes(b"sample-weights")
|
|
return model
|
|
|
|
def test_repairs_list_and_missing_qwen_chat_template(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = self._model(
|
|
root,
|
|
{
|
|
"eos_token": "<|im_end|>",
|
|
"extra_special_tokens": ["<|im_start|>", "<|im_end|>"],
|
|
"additional_special_tokens": ["<|im_start|>"],
|
|
},
|
|
)
|
|
view = Path(prepare_model(source, work_root=root))
|
|
fixed = json.loads((view / "tokenizer_config.json").read_text("utf-8"))
|
|
original = json.loads((source / "tokenizer_config.json").read_text("utf-8"))
|
|
self.assertIn("extra_special_tokens", original)
|
|
self.assertNotIn("extra_special_tokens", fixed)
|
|
self.assertEqual(
|
|
fixed["additional_special_tokens"], ["<|im_start|>", "<|im_end|>"]
|
|
)
|
|
self.assertEqual(fixed["chat_template"], QWEN_CHAT_TEMPLATE)
|
|
self.assertEqual((view / "model.safetensors").read_bytes(), b"sample-weights")
|
|
self.assertTrue((view / ".modelhub-compat-report.json").is_file())
|
|
|
|
def test_no_change_does_not_create_view(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = self._model(
|
|
root,
|
|
{"eos_token": "<|im_end|>", "chat_template": "custom template"},
|
|
)
|
|
self.assertIsNone(prepare_model(source, work_root=root))
|
|
self.assertEqual([p.name for p in root.iterdir()], ["model"])
|
|
|
|
def test_known_tokenizer_class_is_replaced_only_with_tokenizer_json(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = self._model(
|
|
root,
|
|
{"tokenizer_class": "TokenizersBackend", "chat_template": "custom"},
|
|
)
|
|
with self.assertRaisesRegex(CompatibilityError, "tokenizer.json"):
|
|
prepare_model(source, work_root=root)
|
|
(source / "tokenizer.json").write_text("{}", encoding="utf-8")
|
|
view = Path(prepare_model(source, work_root=root))
|
|
fixed = json.loads((view / "tokenizer_config.json").read_text("utf-8"))
|
|
self.assertEqual(fixed["tokenizer_class"], "PreTrainedTokenizerFast")
|
|
|
|
def test_unknown_model_family_does_not_get_qwen_template(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
source = self._model(root, {"eos_token": "<|im_end|>"})
|
|
(source / "config.json").write_text(
|
|
json.dumps({"model_type": "llama"}), encoding="utf-8"
|
|
)
|
|
self.assertIsNone(prepare_model(source, work_root=root))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|