fix

tests/test_library_compatibility.py

@@ -0,0 +1,58 @@
+from pathlib import Path
+
+import pytest
+import torch
+from pydantic import ValidationError
+
+from style_bert_vits2.models.utils.checkpoints import load_checkpoint, save_checkpoint
+from style_bert_vits2.nlp.japanese.user_dict.word_model import UserDictWord
+
+
+def _user_dict_word(**overrides: object) -> UserDictWord:
+    values = {
+        "surface": "test",
+        "priority": 5,
+        "part_of_speech": "名詞",
+        "part_of_speech_detail_1": "一般",
+        "part_of_speech_detail_2": "*",
+        "part_of_speech_detail_3": "*",
+        "inflectional_type": "*",
+        "inflectional_form": "*",
+        "stem": "*",
+        "yomi": "テスト",
+        "pronunciation": "テスト",
+        "accent_type": 1,
+        "accent_associative_rule": "*",
+    }
+    values.update(overrides)
+    return UserDictWord.model_validate(values)
+
+
+def test_user_dict_word_uses_pydantic_v2_validation() -> None:
+    word = _user_dict_word()
+
+    assert word.surface == "ｔｅｓｔ"
+    assert word.mora_count == 3
+
+    with pytest.raises(ValidationError):
+        _user_dict_word(pronunciation="てすと")
+
+
+def test_checkpoint_round_trip_with_safe_torch_load(tmp_path: Path) -> None:
+    checkpoint_path = tmp_path / "checkpoint.pth"
+    source = torch.nn.Linear(2, 1)
+    optimizer = torch.optim.AdamW(source.parameters())
+    save_checkpoint(source, optimizer, 1e-4, 12, checkpoint_path)
+
+    destination = torch.nn.Linear(2, 1)
+    destination_optimizer = torch.optim.AdamW(destination.parameters())
+    _, _, learning_rate, iteration = load_checkpoint(
+        checkpoint_path, destination, destination_optimizer
+    )
+
+    assert learning_rate == 1e-4
+    assert iteration == 12
+    for source_parameter, destination_parameter in zip(
+        source.parameters(), destination.parameters()
+    ):
+        assert torch.equal(source_parameter, destination_parameter)