tests/test_shapes.py

"""Shape and integration tests for all modules."""

import torch
import pytest

from uraionspec.models import (
    DSparkDraftModel, ConfidenceHead, compute_accept_rate,
)
from uraionspec.decoding import (
    verify_block, ThroughputProfile, hardware_aware_prefix_scheduler,
)
from uraionspec.training import compute_dspark_loss, build_position_weights


class TestModelShapes:
    """Test tensor shapes through all model components."""

    @pytest.fixture
    def model(self):
        return DSparkDraftModel(
            vocab_size=100,
            hidden_size=64,
            num_layers=2,
            num_attention_heads=4,
            intermediate_size=128,
            markov_rank=16,
            markov_head_type="vanilla",
            use_confidence_head=True,
        )

    def test_forward_shapes(self, model):
        batch, block_size = 2, 4
        input_ids = torch.randint(0, 100, (batch, block_size))
        prev_ids = torch.randint(0, 100, (batch, block_size))

        output = model(input_ids, prev_token_ids=prev_ids, return_confidence=True)

        assert output["base_logits"].shape == (batch, block_size, 100)
        assert output["hidden_states"].shape == (batch, block_size, 64)
        assert output["draft_logits"].shape == (batch, block_size, 100)
        assert output["confidence_logits"].shape == (batch, block_size)

    def test_forward_no_confidence(self, model):
        model.use_confidence_head = False
        batch, block_size = 2, 4
        input_ids = torch.randint(0, 100, (batch, block_size))

        output = model(input_ids, return_confidence=False)
        assert "confidence_logits" not in output

    def test_sample_block_shapes(self, model):
        batch, gamma = 2, 4
        anchor_ids = torch.tensor([10, 20])

        output = model.sample_block(anchor_ids, gamma=gamma, temperature=0.0)

        assert output["draft_tokens"].shape == (batch, gamma)
        assert output["draft_logits"].shape == (batch, gamma, 100)
        assert output["confidence"].shape == (batch, gamma)


class TestLossShapes:
    """Test loss computation shapes."""

    def test_loss_shapes(self):
        B, N, block_size, V = 2, 3, 4, 100
        draft_logits = torch.randn(B, N, block_size, V, requires_grad=True)
        target_ids = torch.randint(0, V, (B, N, block_size))
        eval_mask = torch.ones(B, N, block_size, dtype=torch.bool)

        result = compute_dspark_loss(
            draft_logits=draft_logits,
            target_ids=target_ids,
            eval_mask=eval_mask,
            ce_alpha=0.1,
            tv_alpha=0.0,
            conf_alpha=0.0,
        )

        loss = result["loss"]
        loss.backward()
        assert loss.item() > 0
        assert draft_logits.grad is not None
        assert draft_logits.grad.shape == (B, N, block_size, V)

    def test_loss_with_all_terms(self):
        B, N, block_size, V = 2, 3, 4, 100
        draft_logits = torch.randn(B, N, block_size, V, requires_grad=True)
        target_ids = torch.randint(0, V, (B, N, block_size))
        target_logits = torch.randn(B, N, block_size, V)
        eval_mask = torch.ones(B, N, block_size, dtype=torch.bool)
        confidence_pred = torch.randn(B, N, block_size, requires_grad=True)
        confidence_targets = torch.sigmoid(torch.randn(B, N, block_size))

        result = compute_dspark_loss(
            draft_logits=draft_logits,
            target_ids=target_ids,
            eval_mask=eval_mask,
            target_logits=target_logits,
            confidence_pred=confidence_pred,
            confidence_targets=confidence_targets,
            ce_alpha=0.1,
            tv_alpha=0.9,
            conf_alpha=1.0,
        )

        result["loss"].backward()
        assert result["loss"].item() > 0

    def test_loss_decay_weights(self):
        weights = build_position_weights(7, 7.0, "cpu")
        assert weights.shape == (7,)
        assert weights[0] > weights[-1]  # Earlier positions have higher weight
        assert (weights > 0).all()


class TestConfidenceHeadShapes:
    """Test confidence head shapes."""

    def test_confidence_head(self):
        batch, gamma, d, r = 2, 4, 64, 16
        head = ConfidenceHead(hidden_size=d, markov_rank=r)

        hidden = torch.randn(batch, gamma, d)
        prev_emb = torch.randn(batch, gamma, r)

        logits = head(hidden, prev_emb)
        assert logits.shape == (batch, gamma)

        probs = torch.sigmoid(logits)
        assert (probs >= 0).all() and (probs <= 1).all()

    def test_accept_rate_shape(self):
        batch, gamma, V = 2, 4, 100
        draft_logits = torch.randn(batch, gamma, V)
        target_logits = torch.randn(batch, gamma, V)

        rate = compute_accept_rate(draft_logits, target_logits)
        assert rate.shape == (batch, gamma)
        assert (rate >= 0).all() and (rate <= 1).all()


class TestIntegration:
    """End-to-end integration tests."""

    def test_draft_verify_cycle(self):
        """Test draft -> verify -> accept cycle end-to-end."""
        V = 50
        model = DSparkDraftModel(
            vocab_size=V, hidden_size=32, num_layers=1,
            num_attention_heads=2, intermediate_size=64,
            markov_rank=8, markov_head_type="vanilla",
            use_confidence_head=True,
        )

        batch, gamma = 2, 4
        anchor_ids = torch.tensor([10, 20])

        # Draft
        draft_out = model.sample_block(anchor_ids, gamma=gamma, temperature=0.0)
        assert draft_out["draft_tokens"].shape == (batch, gamma)
        assert draft_out["draft_logits"].shape == (batch, gamma, V)
        assert draft_out["confidence"].shape == (batch, gamma)

        # Mock target logits (just noise for testing)
        target_logits = torch.randn(batch, gamma, V)

        # Accept
        accepted, num_accepted, bonus = verify_block(
            draft_out["draft_tokens"],
            draft_out["draft_logits"],
            target_logits,
        )
        assert num_accepted.shape == (batch,)
        assert bonus.shape == (batch,)

        # Scheduler
        profile = ThroughputProfile(None)
        conf_list = [draft_out["confidence"][b] for b in range(batch)]
        lengths = hardware_aware_prefix_scheduler(conf_list, profile, gamma)
        assert len(lengths) == batch
        assert all(0 <= length <= gamma for length in lengths)

    def test_train_step_shape(self):
        """Test that all components in a training step have correct shapes."""
        V = 50
        B, N, bs = 2, 2, 4

        draft_logits = torch.randn(B, N, bs, V, requires_grad=True)
        target_ids = torch.randint(0, V, (B, N, bs))
        eval_mask = torch.ones(B, N, bs, dtype=torch.bool)
        target_logits = torch.randn(B, N, bs, V)

        # Confidence targets
        confidence_pred = torch.randn(B, N, bs, requires_grad=True)
        confidence_targets = torch.sigmoid(
            compute_accept_rate(draft_logits, target_logits).detach()
        )

        loss_dict = compute_dspark_loss(
            draft_logits=draft_logits,
            target_ids=target_ids,
            eval_mask=eval_mask,
            target_logits=target_logits,
            confidence_pred=confidence_pred,
            confidence_targets=confidence_targets,
            ce_alpha=0.1,
            tv_alpha=0.9,
            conf_alpha=1.0,
        )

        loss_dict["loss"].backward()
        assert draft_logits.grad is not None
        assert confidence_pred.grad is not None