test_models.py

"""
Test script to verify model implementations work correctly
"""

import sys
import os

# Add current directory to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

def test_imports():
    """Test if all required modules can be imported"""
    try:
        import torch
        print("✓ PyTorch imported successfully")
        
        import numpy as np
        print("✓ NumPy imported successfully")
        
        # Test our model modules
        from taskB import EncoderRNN, DecoderRNN, Seq2seqRNN, HierEncoderRNN, Decoder2RNN
        print("✓ Model classes imported successfully")
        
        return True
    except ImportError as e:
        print(f"✗ Import error: {e}")
        return False

def test_basic_models():
    """Test basic model instantiation"""
    try:
        from taskB import EncoderRNN, DecoderRNN, Seq2seqRNN
        
        vocab_size = 1000
        embed_dim = 300
        hidden_dim = 300
        
        # Test EncoderRNN
        encoder = EncoderRNN(vocab_size, embed_dim, hidden_dim)
        print("✓ EncoderRNN created successfully")
        
        # Test DecoderRNN
        decoder = DecoderRNN(vocab_size, embed_dim, hidden_dim)
        print("✓ DecoderRNN created successfully")
        
        # Test Seq2seqRNN
        model = Seq2seqRNN(vocab_size, embed_dim, hidden_dim)
        print("✓ Seq2seqRNN created successfully")
        
        return True
    except Exception as e:
        print(f"✗ Model creation error: {e}")
        return False

def test_advanced_models():
    """Test advanced model instantiation"""
    try:
        from taskB import HierEncoderRNN, Decoder2RNN, Seq2seqRNN
        
        vocab_size = 1000
        embed_dim = 300
        hidden_dim = 300
        
        # Test HierEncoderRNN
        hier_encoder = HierEncoderRNN(vocab_size, embed_dim, hidden_dim)
        print("✓ HierEncoderRNN created successfully")
        
        # Test Decoder2RNN
        decoder2 = Decoder2RNN(vocab_size, embed_dim, hidden_dim)
        print("✓ Decoder2RNN created successfully")
        
        # Test Seq2seqRNN with advanced components
        model_hier = Seq2seqRNN(vocab_size, embed_dim, hidden_dim, encoder_type='hierarchical')
        print("✓ Seq2seqRNN with hierarchical encoder created successfully")
        
        model_dual = Seq2seqRNN(vocab_size, embed_dim, hidden_dim, decoder_type='dual')
        print("✓ Seq2seqRNN with dual decoder created successfully")
        
        return True
    except Exception as e:
        print(f"✗ Advanced model creation error: {e}")
        return False

def test_forward_pass():
    """Test forward pass with dummy data"""
    try:
        import torch
        from taskB import Seq2seqRNN
        
        vocab_size = 1000
        embed_dim = 300
        hidden_dim = 300
        batch_size = 2
        seq_len = 10
        
        # Create model
        model = Seq2seqRNN(vocab_size, embed_dim, hidden_dim)
        
        # Create dummy data
        src = torch.randint(0, vocab_size, (batch_size, seq_len))
        tgt = torch.randint(0, vocab_size, (batch_size, seq_len))
        src_lengths = [seq_len] * batch_size
        
        # Forward pass
        output = model(src, src_lengths, tgt)
        
        print(f"✓ Forward pass successful, output shape: {output.shape}")
        return True
    except Exception as e:
        print(f"✗ Forward pass error: {e}")
        return False

def main():
    """Run all tests"""
    print("Testing NLP Assignment 2 Models")
    print("=" * 40)
    
    tests = [
        ("Import Test", test_imports),
        ("Basic Models Test", test_basic_models),
        ("Advanced Models Test", test_advanced_models),
        ("Forward Pass Test", test_forward_pass)
    ]
    
    results = []
    for test_name, test_func in tests:
        print(f"\n{test_name}:")
        result = test_func()
        results.append((test_name, result))
    
    print("\n" + "=" * 40)
    print("Test Summary:")
    for test_name, result in results:
        status = "PASS" if result else "FAIL"
        print(f"  {test_name}: {status}")
    
    all_passed = all(result for _, result in results)
    print(f"\nOverall: {'ALL TESTS PASSED' if all_passed else 'SOME TESTS FAILED'}")
    
    return all_passed

if __name__ == "__main__":
    main()