scripts/test_context_encoder.py

"""
Test script for Context Encoder module
"""

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

import torch
from models.context_encoder import ContextEncoder, ContextEncoderWrapper


def test_context_encoder():
    """Test Context Encoder with sample texts"""
    print("=" * 60)
    print("🧪 Testing Context Encoder")
    print("=" * 60)
    
    # Initialize encoder
    device = "cuda" if torch.cuda.is_available() else "cpu"
    encoder = ContextEncoder(
        model_name="all-MiniLM-L6-v2",
        device=device
    )
    
    # Sample texts (simulating session context)
    sample_texts = [
        "User: I just adopted a second dog! He is very energetic.\nAssistant: That's wonderful! What breed is he?",
        "User: I've been working a lot of extra hours lately.\nAssistant: That sounds exhausting. Make sure to take breaks.",
        "User: I love hiking in the mountains on weekends.\nAssistant: That's great exercise! Which trails do you prefer?",
        "User: I'm learning to play the guitar.\nAssistant: That's awesome! How long have you been practicing?"
    ]
    
    print(f"\n📝 Sample texts: {len(sample_texts)}")
    for i, text in enumerate(sample_texts):
        print(f"\n  Text {i+1}: {text[:80]}...")
    
    # Encode texts
    print("\n🔄 Encoding texts...")
    embeddings = encoder(sample_texts)
    
    print(f"\n✅ Encoding complete!")
    print(f"   Output shape: {embeddings.shape}")
    print(f"   Output dtype: {embeddings.dtype}")
    print(f"   Output device: {embeddings.device}")
    
    # Check properties
    print(f"\n🔍 Checking properties:")
    print(f"   - No gradients: {not embeddings.requires_grad}")
    print(f"   - Normalized (L2 norm ≈ 1): {torch.allclose(embeddings.norm(dim=1), torch.ones(len(sample_texts)).to(device), atol=1e-5)}")
    
    # Test similarity (optional)
    print(f"\n📊 Cosine similarities:")
    from sklearn.metrics.pairwise import cosine_similarity
    sim_matrix = cosine_similarity(embeddings.cpu().numpy())
    
    for i in range(len(sample_texts)):
        for j in range(i + 1, len(sample_texts)):
            print(f"   Text {i+1} vs Text {j+1}: {sim_matrix[i, j]:.4f}")
    
    # Test wrapper
    print(f"\n🔄 Testing ContextEncoderWrapper...")
    wrapper = ContextEncoderWrapper(encoder)
    wrapper_embeddings = wrapper(sample_texts)
    
    print(f"   Wrapper output shape: {wrapper_embeddings.shape}")
    print(f"   Wrapper embedding dim: {wrapper.get_embedding_dim()}")
    
    # Verify outputs match
    assert torch.allclose(embeddings, wrapper_embeddings), "Wrapper output doesn't match!"
    print(f"   ✅ Wrapper output matches encoder output")
    
    # Test batch encoding
    print(f"\n🔄 Testing batch encoding (large batch)...")
    large_batch = sample_texts * 10  # 40 texts
    batch_embeddings = encoder.encode_batch(large_batch, batch_size=16)
    
    print(f"   Large batch size: {len(large_batch)}")
    print(f"   Batch encoding output shape: {batch_embeddings.shape}")
    
    print("\n" + "=" * 60)
    print("✅ Context Encoder test completed successfully!")
    print("=" * 60)


if __name__ == "__main__":
    test_context_encoder()