Spaces:
Running
Running
| """Stage 2.3 — Decoder. Hermetic tests on random features + synthetic captions.""" | |
| import pytest | |
| import torch | |
| from capit.config import config | |
| from capit.losses import caption_loss | |
| from capit.models.decoder import Decoder | |
| V = 12 | |
| B = 3 | |
| L = config.encoded_size**2 | |
| # captions sorted by length descending; PAD=0 START=1 END=2 | |
| CAPTIONS = torch.tensor( | |
| [ | |
| [1, 5, 6, 7, 8, 2], | |
| [1, 5, 6, 2, 0, 0], | |
| [1, 7, 2, 0, 0, 0], | |
| ] | |
| ) | |
| LENGTHS = torch.tensor([6, 4, 3]) | |
| def _seed(): | |
| torch.manual_seed(config.seed) | |
| def _features() -> torch.Tensor: | |
| return torch.randn(B, L, config.encoder_dim) | |
| def test_forward_shapes(): | |
| logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS) | |
| assert decode_lengths == [5, 3, 2] | |
| assert logits.shape == (B, 5, V) | |
| assert alphas.shape == (B, 5, L) | |
| def test_no_nan(): | |
| logits, alphas, _ = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS) | |
| assert not torch.isnan(logits).any() | |
| assert not torch.isnan(alphas).any() | |
| def test_shrinking_leaves_finished_rows_zero(): | |
| logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS) | |
| for i, dl in enumerate(decode_lengths): | |
| assert torch.all(logits[i, dl:] == 0) | |
| assert torch.all(alphas[i, dl:] == 0) | |
| def test_padded_positions_contribute_zero_loss(): | |
| logits, alphas, decode_lengths = Decoder(vocab_size=V)(_features(), CAPTIONS, LENGTHS) | |
| base = caption_loss(logits, alphas, CAPTIONS) | |
| corrupted = logits.clone() | |
| for i, dl in enumerate(decode_lengths): | |
| corrupted[i, dl:] = 999.0 | |
| assert torch.isclose(base, caption_loss(corrupted, alphas, CAPTIONS)) | |
| def test_forward_rejects_unsorted_batch(): | |
| captions = torch.tensor([[1, 7, 2, 0, 0, 0], [1, 5, 6, 7, 8, 2]]) # ascending lengths | |
| lengths = torch.tensor([3, 6]) | |
| with pytest.raises(ValueError): | |
| Decoder(vocab_size=V)(torch.randn(2, L, config.encoder_dim), captions, lengths) | |
| def test_gradient_reaches_decoder_params(): | |
| dec = Decoder(vocab_size=V) | |
| logits, alphas, _ = dec(_features(), CAPTIONS, LENGTHS) | |
| caption_loss(logits, alphas, CAPTIONS).backward() | |
| for p in (dec.embedding.weight, dec.fc.weight, dec.f_beta.weight): | |
| assert p.grad is not None and p.grad.abs().sum() > 0 | |
| def test_greedy_returns_tokens_and_restores_mode(): | |
| dec = Decoder(vocab_size=V) | |
| dec.train() | |
| out = dec.greedy(_features(), start_id=1, end_id=2, max_len=10) | |
| assert len(out) == B | |
| assert all(isinstance(seq, list) and 2 not in seq for seq in out) | |
| assert dec.training | |