test_inference.py

#!/usr/bin/env python3
"""
Test script for TransPolymer inference
Tests the inference logic locally before SageMaker deployment
"""

import json
import os
import sys
import torch
from unittest.mock import Mock
import traceback

def test_local_inference():
    """
    Test the inference script locally
    """
    print("🧪 Testing TransPolymer inference locally...")
    
    try:
        # Import the inference module
        from inference import ModelHandler, model_fn, input_fn, predict_fn, output_fn
        
        print("✅ Successfully imported inference module")
        
        # Test 1: Model initialization
        print("\n1️⃣ Testing model initialization...")
        
        # Create mock context
        class MockContext:
            def __init__(self):
                self.system_properties = {"model_dir": "./ckpt/pretrain.pt"}
        
        context = MockContext()
        
        # Initialize model handler
        handler = ModelHandler()
        try:
            handler.initialize(context)
            print("✅ Model initialization successful")
        except Exception as e:
            print(f"⚠️ Model initialization warning: {e}")
            print("This is expected if model files are not available")
        
        # Test 2: Input processing
        print("\n2️⃣ Testing input processing...")
        
        test_inputs = [
            {
                "smiles": "CC(C)(C)OC(=O)NC1=CC=CC=C1",
                "property": "conductivity",
                "model_type": "PE_I"
            },
            {
                "smiles": "c1ccc2c(c1)oc1ccccc12",
                "property": "band_gap", 
                "model_type": "Egc"
            }
        ]
        
        # Test input_fn
        request_body = json.dumps(test_inputs)
        parsed_input = input_fn(request_body, "application/json")
        print(f"✅ Input parsing successful: {len(parsed_input)} items")
        
        # Test single input
        single_input = json.dumps(test_inputs[0])
        single_parsed = input_fn(single_input, "application/json")
        print(f"✅ Single input parsing successful: {len(single_parsed)} items")
        
        # Test 3: Output formatting
        print("\n3️⃣ Testing output formatting...")
        
        mock_predictions = [
            {
                "prediction": 1.23e-4,
                "smiles": "CC(C)(C)OC(=O)NC1=CC=CC=C1",
                "property": "conductivity",
                "model_type": "PE_I",
                "units": "S/cm",
                "description": "Polymer Electrolyte Conductivity",
                "confidence": 0.95
            }
        ]
        
        output = output_fn(mock_predictions, "application/json")
        parsed_output = json.loads(output)
        print(f"✅ Output formatting successful: {parsed_output}")
        
        print("\n4️⃣ Testing error handling...")
        
        # Test invalid content type
        try:
            input_fn("test", "text/plain")
            print("❌ Should have raised error for invalid content type")
        except ValueError:
            print("✅ Correctly handled invalid content type")
        
        # Test empty SMILES
        try:
            empty_input = json.dumps([{"smiles": "", "property": "test"}])
            parsed = input_fn(empty_input, "application/json")
            if handler.tokenizer:  # Only test if tokenizer is available
                handler.preprocess(parsed)
                print("❌ Should have raised error for empty SMILES")
        except (ValueError, AttributeError):
            print("✅ Correctly handled empty SMILES")
        
        print("\n✅ All local tests passed!")
        return True
        
    except ImportError as e:
        print(f"❌ Import error: {e}")
        return False
    except Exception as e:
        print(f"❌ Test failed: {e}")
        traceback.print_exc()
        return False

def test_dependencies():
    """
    Test if all required dependencies are available
    """
    print("📦 Testing dependencies...")
    
    dependencies = [
        ('torch', 'PyTorch'),
        ('transformers', 'Transformers'),
        ('numpy', 'NumPy'),
        ('json', 'JSON (built-in)'),
        ('os', 'OS (built-in)')
    ]
    
    missing = []
    for module, name in dependencies:
        try:
            __import__(module)
            print(f"✅ {name}")
        except ImportError:
            print(f"❌ {name} - Missing!")
            missing.append(name)
    
    # Test optional dependencies
    optional_deps = [
        ('rdkit', 'RDKit'),
        ('sklearn', 'scikit-learn'),
        ('boto3', 'Boto3'),
        ('sagemaker', 'SageMaker Python SDK')
    ]
    
    print("\nOptional dependencies:")
    for module, name in optional_deps:
        try:
            __import__(module)
            print(f"✅ {name}")
        except ImportError:
            print(f"⚠️ {name} - Not available (optional)")
    
    if missing:
        print(f"\n❌ Missing required dependencies: {missing}")
        return False
    else:
        print("\n✅ All required dependencies available!")
        return True

def test_model_files():
    """
    Check if model files exist
    """
    print("📁 Checking model files...")
    
    files_to_check = [
        'PolymerSmilesTokenization.py',
        'inference.py',
        'requirements.txt',
        'ckpt/pretrain.pt/config.json',
        'ckpt/pretrain.pt/pytorch_model.bin'
    ]
    
    for file_path in files_to_check:
        if os.path.exists(file_path):
            size = os.path.getsize(file_path)
            print(f"✅ {file_path} ({size} bytes)")
        else:
            print(f"⚠️ {file_path} - Not found")
    
    return True

def generate_sample_requests():
    """
    Generate sample API requests for testing
    """
    print("\n📝 Sample API requests:")
    
    samples = [
        {
            "description": "Polymer Electrolyte Conductivity",
            "request": {
                "smiles": "CC(C)(C)OC(=O)NC1=CC=CC=C1",
                "property": "conductivity",
                "model_type": "PE_I"
            }
        },
        {
            "description": "Band Gap Prediction",
            "request": {
                "smiles": "c1ccc2c(c1)oc1ccccc12",
                "property": "band_gap",
                "model_type": "Egc"
            }
        },
        {
            "description": "OPV Efficiency",
            "request": {
                "smiles": "c1ccc(cc1)c2ccc(cc2)C3=CC=C(C=C3)c4ccccc4",
                "property": "efficiency",
                "model_type": "OPV"
            }
        }
    ]
    
    for i, sample in enumerate(samples, 1):
        print(f"\n{i}. {sample['description']}:")
        print(f"   curl -X POST https://your-endpoint/invocations \\")
        print(f"        -H 'Content-Type: application/json' \\")
        print(f"        -d '{json.dumps(sample['request'])}'")

def main():
    """
    Main test function
    """
    print("🚀 TransPolymer Inference Testing Suite")
    print("=" * 50)
    
    success = True
    
    # Test 1: Dependencies
    if not test_dependencies():
        success = False
    
    print("\n" + "=" * 50)
    
    # Test 2: Model files
    test_model_files()
    
    print("\n" + "=" * 50)
    
    # Test 3: Local inference
    if not test_local_inference():
        success = False
    
    print("\n" + "=" * 50)
    
    # Generate sample requests
    generate_sample_requests()
    
    print("\n" + "=" * 50)
    
    if success:
        print("🎉 All tests completed successfully!")
        print("\nNext steps:")
        print("1. Ensure you have proper AWS credentials configured")
        print("2. Run: python deploy_sagemaker.py")
        print("3. Test the deployed endpoint with the sample requests above")
    else:
        print("❌ Some tests failed. Please fix the issues before deploying.")
        sys.exit(1)

if __name__ == "__main__":
    main()