| |
|
|
| import pytest |
| import torch |
| import torch.nn as nn |
| from unittest.mock import MagicMock, patch, PropertyMock |
|
|
| from config import get_all_special_tokens, CODEBOOK_SIZE, NUM_SPEECH_TOKENS |
|
|
|
|
| class MockEmbedding(nn.Embedding): |
| def __init__(self, vocab_size=50280, embed_dim=2048): |
| super().__init__(vocab_size, embed_dim) |
|
|
|
|
| class MockLinear(nn.Linear): |
| def __init__(self, in_features=2048, out_features=50280): |
| super().__init__(in_features, out_features, bias=False) |
|
|
|
|
| class TestTokenExpansion: |
| def test_special_token_count(self): |
| import config |
| from config import NUM_SPEECH_TOKENS |
| tokens = get_all_special_tokens() |
| n_plan = len(config.get_plan_tokens()) if config.ENABLE_PLAN_TOKENS else 0 |
| assert len(tokens) == 22 + n_plan + NUM_SPEECH_TOKENS |
| if not config.ENABLE_PLAN_TOKENS: |
| assert len(tokens) == 22 + NUM_SPEECH_TOKENS |
|
|
| def test_speech_tokens_in_list(self): |
| tokens = get_all_special_tokens() |
| assert "[SPEECH_0]" in tokens |
| assert f"[SPEECH_{NUM_SPEECH_TOKENS - 1}]" in tokens |
| assert "[SPEECH_512]" in tokens |
|
|
| def test_control_tokens_in_list(self): |
| tokens = get_all_special_tokens() |
| assert "<THINK>" in tokens |
| assert "</THINK>" in tokens |
| assert "<AUDIO>" in tokens |
| assert "</AUDIO>" in tokens |
| assert "[USER_PROMPT]" in tokens |
|
|
| def test_emotion_tokens_in_list(self): |
| tokens = get_all_special_tokens() |
| assert "[EMO:sarcastic]" in tokens |
| assert "[EMO:neutral]" in tokens |
|
|
| def test_prosody_tokens_in_list(self): |
| tokens = get_all_special_tokens() |
| assert "[PACE:slow]" in tokens |
| assert "[PITCH:high]" in tokens |
| assert "[BREAK:long]" in tokens |
|
|
|
|
| class TestEmbeddingResize: |
| def test_resize_expands_correctly(self): |
| old_size = 50280 |
| new_size = old_size + 1045 |
| embed_dim = 128 |
|
|
| old_embedding = nn.Embedding(old_size, embed_dim) |
| new_embedding = nn.Embedding(new_size, embed_dim) |
|
|
| with torch.no_grad(): |
| new_embedding.weight[:old_size] = old_embedding.weight[:old_size] |
|
|
| assert new_embedding.weight.shape[0] == new_size |
| assert new_embedding.weight.shape[1] == embed_dim |
|
|
| assert torch.allclose( |
| new_embedding.weight[:old_size], |
| old_embedding.weight[:old_size], |
| ) |
|
|
| def test_new_rows_initialized(self): |
| old_size = 100 |
| new_size = 150 |
| embed_dim = 64 |
|
|
| old_embedding = nn.Embedding(old_size, embed_dim) |
| new_embedding = nn.Embedding(new_size, embed_dim) |
|
|
| with torch.no_grad(): |
| new_embedding.weight[:old_size] = old_embedding.weight[:old_size] |
| mean_embed = old_embedding.weight.mean(dim=0) |
| noise = torch.randn(new_size - old_size, embed_dim) * 0.02 |
| new_embedding.weight[old_size:] = mean_embed + noise |
|
|
| new_rows = new_embedding.weight[old_size:] |
| assert not torch.all(new_rows == 0) |
|
|
| def test_lm_head_resize(self): |
| old_size = 50280 |
| new_size = old_size + 1045 |
| hidden_dim = 128 |
|
|
| old_head = nn.Linear(hidden_dim, old_size, bias=False) |
| new_head = nn.Linear(hidden_dim, new_size, bias=False) |
|
|
| with torch.no_grad(): |
| new_head.weight[:old_size] = old_head.weight[:old_size] |
|
|
| assert new_head.weight.shape[0] == new_size |
| assert new_head.weight.shape[1] == hidden_dim |
|
|
|
|
| class TestForwardPassLogic: |
| def test_causal_lm_loss_shape(self): |
| batch_size, seq_len, vocab_size = 2, 10, 100 |
| logits = torch.randn(batch_size, seq_len, vocab_size) |
| labels = torch.randint(0, vocab_size, (batch_size, seq_len)) |
|
|
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
|
|
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100) |
| loss = loss_fct( |
| shift_logits.view(-1, shift_logits.size(-1)), |
| shift_labels.view(-1), |
| ) |
|
|
| assert loss.dim() == 0 |
| assert not torch.isnan(loss) |
|
|
| def test_ignore_index_masks_correctly(self): |
| vocab_size = 100 |
| logits = torch.randn(1, 5, vocab_size) |
| labels = torch.tensor([[1, -100, -100, 42, 7]]) |
|
|
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
|
|
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction="none") |
| per_token = loss_fct( |
| shift_logits.view(-1, shift_logits.size(-1)), |
| shift_labels.view(-1), |
| ) |
|
|
| assert per_token[0] == 0.0 |
| assert per_token[1] == 0.0 |
| assert per_token[2] > 0.0 |
| assert per_token[3] > 0.0 |
|
|
| def test_weighted_loss(self): |
| vocab_size = 100 |
| logits = torch.randn(1, 5, vocab_size) |
| labels = torch.randint(0, vocab_size, (1, 5)) |
| weights = torch.tensor([[1.0, 0.3, 0.3, 1.0, 1.0]]) |
|
|
| shift_logits = logits[..., :-1, :].contiguous() |
| shift_labels = labels[..., 1:].contiguous() |
| shift_weights = weights[..., 1:].contiguous() |
|
|
| loss_fct = nn.CrossEntropyLoss(ignore_index=-100, reduction="none") |
| per_token = loss_fct( |
| shift_logits.view(-1, shift_logits.size(-1)), |
| shift_labels.view(-1), |
| ) |
|
|
| weighted = per_token * shift_weights.view(-1) |
| loss = weighted.sum() / (shift_labels.view(-1) != -100).float().sum().clamp(min=1) |
|
|
| assert loss.dim() == 0 |
| assert not torch.isnan(loss) |
|
|
|
|
| class TestFreezeUnfreeze: |
| def test_freeze_sets_requires_grad_false(self): |
| model = nn.Sequential( |
| nn.Linear(10, 20), |
| nn.Linear(20, 10), |
| ) |
| for param in model.parameters(): |
| param.requires_grad = False |
|
|
| trainable = sum(1 for p in model.parameters() if p.requires_grad) |
| assert trainable == 0 |
|
|
| def test_unfreeze_sets_requires_grad_true(self): |
| model = nn.Sequential( |
| nn.Linear(10, 20), |
| nn.Linear(20, 10), |
| ) |
| for param in model.parameters(): |
| param.requires_grad = False |
|
|
| for param in model.parameters(): |
| param.requires_grad = True |
|
|
| trainable = sum(1 for p in model.parameters() if p.requires_grad) |
| total = sum(1 for _ in model.parameters()) |
| assert trainable == total |
|
|
