tests/test_thinking_dial_integration.py

"""Thinking Dial integration test - verifies N10/N11/N12/N17/N18/N19/N20 fixes."""
import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import pytest
from models.fusion_model import FusionConfig, FusionModel
from models.thinking_dial import ThinkingDialModel, GRPOTrainer, ThinkingConfig


@pytest.fixture
def setup():
    config = FusionConfig(
        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
        block_size=8, latent_dim=16, window_size=64,
    )
    base_model = FusionModel(config)
    base_model.eval()
    td_model = ThinkingDialModel(base_model, ThinkingConfig())
    td_model.eval()
    input_ids = torch.randint(0, 1000, (1, 5))
    return td_model, base_model, input_ids


def test_thinking_dial_model_generate(setup):
    """N11: ThinkingDialModel.generate() exists and works."""
    td_model, _, input_ids = setup
    with torch.no_grad():
        out = td_model.generate(input_ids=input_ids, max_new_tokens=5, thinking_depth=None)
    assert out.shape[0] == 1


def test_thinking_depth_bias_applied(setup):
    """N10/N18: thinking_depth bias produces different logits for different depths."""
    td_model, base_model, input_ids = setup
    with torch.no_grad():
        # Get raw logits without thinking bias
        raw_out = base_model(input_ids, return_dict=True)
        raw_logits_d0 = raw_out.logits[:, -1, :]  # (1, vocab)
        
        # Get logits with depth=0 bias via ThinkingDialModel.generate's hook mechanism
        hook_d0 = ThinkingDialModel._build_thinking_logits_hook(
            0, 1, input_ids.device, td_model.thinking_config,
            td_model.thinking_embedding, td_model.thinking_gate, base_model.lm_head,
        )
        hook_d3 = ThinkingDialModel._build_thinking_logits_hook(
            3, 1, input_ids.device, td_model.thinking_config,
            td_model.thinking_embedding, td_model.thinking_gate, base_model.lm_head,
        )
        biased_logits_d0 = hook_d0(raw_logits_d0.unsqueeze(1)).squeeze(1)
        biased_logits_d3 = hook_d3(raw_logits_d0.unsqueeze(1)).squeeze(1)
    
    # Raw logits should differ from biased logits
    assert not torch.allclose(raw_logits_d0, biased_logits_d0), "Depth=0 bias should change logits"
    # Different depths should produce different logits (bias vectors differ)
    assert not torch.allclose(biased_logits_d0, biased_logits_d3), "Depth=0 and depth=3 should differ"


def test_thinking_dial_n19_single_source(setup):
    """N19: _build_thinking_logits_hook is single source of truth used by both paths."""
    td_model, base_model, input_ids = setup
    # Verify the static method exists
    assert hasattr(ThinkingDialModel, '_build_thinking_logits_hook')
    # Verify it returns None when depth is None
    hook = ThinkingDialModel._build_thinking_logits_hook(
        None, 1, input_ids.device, td_model.thinking_config,
        td_model.thinking_embedding, td_model.thinking_gate, base_model.lm_head,
    )
    assert hook is None


def test_n20_forward_no_thinking_depth(setup):
    """N20: FusionModel.forward() should NOT accept thinking_depth (dead param removed)."""
    _, base_model, input_ids = setup
    import inspect
    sig = inspect.signature(base_model.forward)
    assert 'thinking_depth' not in sig.parameters, \
        "thinking_depth should be removed from FusionModel.forward() — use logits_hook instead"


def test_grpo_trainer_generate_with_thinking(setup):
    """N12: GRPOTrainer.generate_with_thinking() passes depth."""
    td_model, _, input_ids = setup
    trainer = GRPOTrainer(td_model)
    with torch.no_grad():
        texts = trainer.generate_with_thinking(input_ids, thinking_depth=2, max_new_tokens=5)
    assert len(texts) == 1


def test_n18_first_token_has_bias(setup):
    """N18: First sampled token in generate_samples should have thinking bias applied."""
    td_model, _, input_ids = setup
    trainer = GRPOTrainer(td_model)
    torch.manual_seed(42)
    ids_d0, _ = trainer.generate_samples(input_ids, num_samples=1, thinking_depth=0, max_new_tokens=3)
    torch.manual_seed(42)
    ids_d3, _ = trainer.generate_samples(input_ids, num_samples=1, thinking_depth=3, max_new_tokens=3)
    # Same seed but different depth → first generated token should differ (due to bias)
    # Check prompt portion matches
    assert torch.equal(ids_d0[0, :input_ids.shape[1]], ids_d3[0, :input_ids.shape[1]]), \
        "Prompt portion should be identical"
    # First generated token (position after prompt) should differ
    first_gen_0 = ids_d0[0, input_ids.shape[1]].item()
    first_gen_3 = ids_d3[0, input_ids.shape[1]].item()
    # With random weights the biases will differ, so tokens should differ
    # (Not guaranteed for every seed, but highly likely with different bias vectors)
    if first_gen_0 == first_gen_3:
        # Accept if the first token matches but verify the hook was actually applied
        # by checking logits directly
        with torch.no_grad():
            raw = td_model.base_model(input_ids, return_dict=True)
            raw_logits = raw.logits[:, -1, :]
            hook_d0 = ThinkingDialModel._build_thinking_logits_hook(
                0, 1, input_ids.device, td_model.thinking_config,
                td_model.thinking_embedding, td_model.thinking_gate, td_model.base_model.lm_head,
            )
            hook_d3 = ThinkingDialModel._build_thinking_logits_hook(
                3, 1, input_ids.device, td_model.thinking_config,
                td_model.thinking_embedding, td_model.thinking_gate, td_model.base_model.lm_head,
            )
            b0 = hook_d0(raw_logits.unsqueeze(1)).squeeze(1)
            b3 = hook_d3(raw_logits.unsqueeze(1)).squeeze(1)
            assert not torch.allclose(b0, b3), "Hooks must produce different logits"
    else:
        pass  # Best case: tokens differ


def test_generate_samples_with_depth(setup):
    """N12/N17: generate_samples passes thinking_depth, uses KV cache reuse."""
    td_model, _, input_ids = setup
    trainer = GRPOTrainer(td_model)
    with torch.no_grad():
        ids, texts = trainer.generate_samples(
            input_ids, num_samples=2, thinking_depth=1, max_new_tokens=3,
        )
    assert ids.shape[0] >= 2


def test_n23_prompt_mask_correct_dimension():
    """N23: Prompt mask should zero per-token positions, not per-sequence."""
    config = FusionConfig(
        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
        block_size=8, latent_dim=16, window_size=64,
    )
    base_model = FusionModel(config)
    trainer = GRPOTrainer(base_model)
    
    # Test _normalize_logits_to_log_probs with per_token=True returns 2D
    B, L, V = 4, 10, config.vocab_size
    logits = torch.randn(B, L, V)
    labels = torch.randint(1, V, (B, L))  # non-zero labels
    
    per_token = trainer._normalize_logits_to_log_probs(logits, labels, per_token=True)
    per_seq = trainer._normalize_logits_to_log_probs(logits, labels, per_token=False)
    
    assert per_token.dim() == 2 and per_token.shape == (B, L-1), f"Expected ({B}, {L-1}), got {per_token.shape}"
    assert per_seq.dim() == 1 and per_seq.shape == (B,), f"Expected ({B},), got {per_seq.shape}"
    
    # Verify per_seq = per_token.sum(dim=1)
    assert torch.allclose(per_seq, per_token.sum(dim=1), atol=1e-5), "per_seq should equal per_token sum"
    
    # Verify masking: zero out first 3 positions, sum should differ
    masked = per_token.clone()
    masked[:, :3] = 0.0
    assert not torch.allclose(masked.sum(dim=1), per_seq), "Masking should change the sum"


def test_n24_pure_fusion_model_no_crash():
    """N24: GRPOTrainer with pure FusionModel should not crash (hasattr guard)."""
    config = FusionConfig(
        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
        block_size=8, latent_dim=16, window_size=64,
    )
    base_model = FusionModel(config)
    base_model.eval()
    trainer = GRPOTrainer(base_model)
    input_ids = torch.randint(0, 1000, (1, 5))
    
    # Should not raise AttributeError even with thinking_depth=None
    with torch.no_grad():
        ids, texts = trainer.generate_samples(
            input_ids, num_samples=1, thinking_depth=None, max_new_tokens=3,
        )
    assert ids.shape[0] == 1


def test_m7_forward_uses_shared_hook(setup):
    """M7: ThinkingDialModel.forward() uses _build_thinking_logits_hook (single source)."""
    td_model, _, input_ids = setup
    with torch.no_grad():
        out_no_depth = td_model(input_ids)
        out_depth0 = td_model(input_ids, thinking_depth=0)
        out_depth3 = td_model(input_ids, thinking_depth=3)
    
    # Different depths should produce different logits
    assert not torch.allclose(out_depth0.logits, out_depth3.logits), \
        "Different thinking depths should produce different logits via forward()"
    # depth=0 should still differ from no depth
    assert not torch.allclose(out_no_depth.logits, out_depth0.logits), \
        "No depth vs depth=0 should differ"


def test_s2_train_step_accepts_thinking_depth(setup):
    """S2: train_step() accepts thinking_depth parameter."""
    td_model, _, input_ids = setup
    trainer = GRPOTrainer(td_model)
    # Verify the parameter exists
    import inspect
    sig = inspect.signature(trainer.train_step)
    assert 'thinking_depth' in sig.parameters, "train_step should accept thinking_depth"


def test_n25_mask_start_correct():
    """N25: mask_start should be prompt_len - 2 (first gen token at index prompt_len-2)."""
    config = FusionConfig(
        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
        block_size=8, latent_dim=16, window_size=64,
    )
    base_model = FusionModel(config)
    trainer = GRPOTrainer(base_model)
    
    # Test the mask_start calculation: prompt_len=5 -> mask_start=3
    # This means index 0,1,2 (prompt tokens) are zeroed, index 3+ (gen tokens) are kept
    prompt_len = 5
    mask_start = max(prompt_len - 2, 0)
    assert mask_start == 3, f"Expected mask_start=3, got {mask_start}"
    
    # Verify: index 3 is first gen token (not zeroed)
    assert mask_start == 3, "First gen token should NOT be masked out"


def test_reward_fn_string_safety():
    """S2 FIX: reward_fn as string should not crash compute_reward."""
    config = FusionConfig(
        vocab_size=1000, hidden_size=256, num_hidden_layers=2,
        num_attention_heads=4, num_key_value_heads=2, intermediate_size=512,
        block_size=8, latent_dim=16, window_size=64,
    )
    base_model = FusionModel(config)
    # Register a test reward function
    GRPOTrainer.register_reward_fn('test_reward', lambda p, r: 1.0)
    
    trainer = GRPOTrainer(base_model, reward_fn='test_reward')  # Set as string
    # Should not raise TypeError
    reward = trainer.compute_reward('prompt', 'response')
    assert reward == 1.0, "String reward_fn should be looked up from registry"