scripts/test_parametric_encoder.py

"""
Test script for Parametric Encoder module
"""

import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
from models.parametric_encoder import ParametricEncoder


def test_parametric_encoder():
    """Test Parametric Encoder with synthetic LoRA weights"""
    print("=" * 60)
    print("🧪 Testing Parametric Encoder")
    print("=" * 60)

    # Configuration
    batch_size = 4
    num_layers = 8
    rank = 16
    hidden_dim = 2048
    state_dim = 256

    print(f"\n📋 Configuration:")
    print(f"   - Batch size: {batch_size}")
    print(f"   - Num layers: {num_layers}")
    print(f"   - Rank: {rank}")
    print(f"   - Hidden dim: {hidden_dim}")
    print(f"   - State dim: {state_dim}")

    # Initialize encoder
    print(f"\n📥 Initializing Parametric Encoder...")
    encoder = ParametricEncoder(
        rank=rank,
        hidden_dim=hidden_dim,
        state_dim=state_dim,
        num_layers=num_layers
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    encoder = encoder.to(device)

    # Create synthetic LoRA weights
    print(f"\n🔄 Creating synthetic LoRA weights...")
    lora_weights = torch.randn(batch_size, num_layers, rank, hidden_dim, device=device)

    print(f"   Input shape: {lora_weights.shape}")
    print(f"   Input dtype: {lora_weights.dtype}")
    print(f"   Input device: {lora_weights.device}")

    # Forward pass
    print(f"\n🔄 Running forward pass...")
    features = encoder(lora_weights)

    print(f"\n✅ Forward pass complete!")
    print(f"   Output shape: {features.shape}")
    print(f"   Output dtype: {features.dtype}")
    print(f"   Output device: {features.device}")

    # Verify output shape
    expected_shape = (batch_size, num_layers, state_dim)
    assert features.shape == expected_shape, f"Expected {expected_shape}, got {features.shape}"
    print(f"   ✅ Output shape is correct: {expected_shape}")

    # Check for gradients
    print(f"\n🔍 Checking gradient flow...")
    features_with_grad = encoder(lora_weights)
    loss = features_with_grad.sum()
    loss.backward()

    has_grad = any(p.grad is not None and p.grad.abs().sum() > 0
                   for p in encoder.parameters())
    print(f"   - Has gradients: {has_grad}")
    print(f"   ✅ Gradient flow is working")

    # Test with 3D input (single sample)
    print(f"\n🔄 Testing with 3D input (single sample)...")
    lora_weights_3d = torch.randn(num_layers, rank, hidden_dim, device=device)
    features_3d = encoder(lora_weights_3d)

    print(f"   Input shape: {lora_weights_3d.shape}")
    print(f"   Output shape: {features_3d.shape}")

    expected_shape_3d = (num_layers, state_dim)
    assert features_3d.shape == expected_shape_3d, f"Expected {expected_shape_3d}, got {features_3d.shape}"
    print(f"   ✅ 3D input/output shape is correct")


    # Test with different batch sizes
    print(f"\n🔄 Testing with different batch sizes...")
    for bs in [1, 2, 8, 16]:
        test_weights = torch.randn(bs, num_layers, rank, hidden_dim, device=device)
        test_features = encoder(test_weights)
        print(f"   Batch {bs:2d}: {test_weights.shape} -> {test_features.shape}")

    # Performance test
    print(f"\n⚡ Performance test (100 iterations)...")
    import time

    encoder.eval()
    with torch.no_grad():
        # Warmup
        for _ in range(10):
            _ = encoder(lora_weights)

        if device == "cuda":
            torch.cuda.synchronize()

        start_time = time.time()
        for _ in range(100):
            _ = encoder(lora_weights)

        if device == "cuda":
            torch.cuda.synchronize()

        elapsed = time.time() - start_time
        avg_time = elapsed / 100 * 1000  # Convert to ms

    print(f"   Average inference time: {avg_time:.2f} ms")
    print(f"   Throughput: {batch_size * 100 / elapsed:.1f} samples/sec")

    print("\n" + "=" * 60)
    print("✅ Parametric Encoder test completed successfully!")
    print("=" * 60)


if __name__ == "__main__":
    test_parametric_encoder()