59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
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)
|