107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple test script to verify Vertex AI multimodal embeddings work correctly.
|
|
Run this script to test the embedding generation before deploying.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import logging
|
|
from PIL import Image
|
|
import io
|
|
import numpy as np
|
|
|
|
# Add the current directory to the path so we can import main
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def create_test_image():
|
|
"""Create a simple test image"""
|
|
# Create a simple 100x100 RGB image with a red square
|
|
img = Image.new('RGB', (100, 100), color='white')
|
|
pixels = img.load()
|
|
|
|
# Draw a red square in the center
|
|
for i in range(25, 75):
|
|
for j in range(25, 75):
|
|
pixels[i, j] = (255, 0, 0) # Red
|
|
|
|
# Convert to bytes
|
|
img_bytes = io.BytesIO()
|
|
img.save(img_bytes, format='PNG')
|
|
return img_bytes.getvalue()
|
|
|
|
def test_vertex_ai_embeddings():
|
|
"""Test Vertex AI embedding generation"""
|
|
try:
|
|
# Set required environment variables for testing
|
|
os.environ['GOOGLE_CLOUD_PROJECT'] = os.environ.get('GOOGLE_CLOUD_PROJECT', 'your-project-id')
|
|
os.environ['VERTEX_AI_LOCATION'] = os.environ.get('VERTEX_AI_LOCATION', 'us-central1')
|
|
|
|
# Import the function we want to test
|
|
from main import generate_image_embeddings
|
|
|
|
# Create test image
|
|
logger.info("Creating test image...")
|
|
test_image_data = create_test_image()
|
|
logger.info(f"Created test image with {len(test_image_data)} bytes")
|
|
|
|
# Generate embeddings
|
|
logger.info("Generating embeddings using Vertex AI...")
|
|
embeddings = generate_image_embeddings(test_image_data)
|
|
|
|
if embeddings is None:
|
|
logger.error("Failed to generate embeddings!")
|
|
return False
|
|
|
|
# Verify embeddings
|
|
logger.info(f"Generated embeddings with shape: {embeddings.shape}")
|
|
logger.info(f"Embeddings dtype: {embeddings.dtype}")
|
|
logger.info(f"Embeddings range: [{embeddings.min():.4f}, {embeddings.max():.4f}]")
|
|
logger.info(f"Embeddings norm: {np.linalg.norm(embeddings):.4f}")
|
|
|
|
# Basic validation
|
|
assert isinstance(embeddings, np.ndarray), "Embeddings should be numpy array"
|
|
assert embeddings.dtype == np.float32, "Embeddings should be float32"
|
|
assert len(embeddings.shape) == 1, "Embeddings should be 1D array"
|
|
assert embeddings.shape[0] == 1408, f"Expected 1408 dimensions, got {embeddings.shape[0]}"
|
|
|
|
# Check if normalized (should be close to 1.0)
|
|
norm = np.linalg.norm(embeddings)
|
|
assert 0.9 <= norm <= 1.1, f"Embeddings should be normalized, norm is {norm}"
|
|
|
|
logger.info("✅ All tests passed! Vertex AI embeddings are working correctly.")
|
|
return True
|
|
|
|
except ImportError as e:
|
|
logger.error(f"Import error: {e}")
|
|
logger.error("Make sure you have installed the required dependencies:")
|
|
logger.error("pip install google-cloud-aiplatform")
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Test failed with error: {e}")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
logger.info("Testing Vertex AI multimodal embeddings...")
|
|
|
|
# Check if required environment variables are set
|
|
project_id = os.environ.get('GOOGLE_CLOUD_PROJECT')
|
|
if not project_id:
|
|
logger.error("Please set GOOGLE_CLOUD_PROJECT environment variable")
|
|
logger.error("Example: export GOOGLE_CLOUD_PROJECT=your-project-id")
|
|
sys.exit(1)
|
|
|
|
logger.info(f"Using project: {project_id}")
|
|
|
|
success = test_vertex_ai_embeddings()
|
|
|
|
if success:
|
|
logger.info("🎉 Test completed successfully!")
|
|
sys.exit(0)
|
|
else:
|
|
logger.error("❌ Test failed!")
|
|
sys.exit(1) |