nomic / test_local.py
AkinyemiAra's picture
Upload 10 files
67e79b5 verified
import requests
from PIL import Image
import io
import sys
import json
def test_local_server(image_path=None):
"""
Test the local Gradio server by sending a request to generate embeddings for an image
Args:
image_path: Path to the image file. If None, a test URL will be used.
"""
# Local server URL (default Gradio port)
server_url = "http://localhost:7860/api/predict"
if image_path:
try:
# Load the image
image = Image.open(image_path)
# Convert image to bytes
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format=image.format if image.format else 'PNG')
img_byte_arr = img_byte_arr.getvalue()
# Prepare the request
files = {
'data': ('image.png', img_byte_arr, 'image/png')
}
# Send the request
response = requests.post(server_url, files=files)
except Exception as e:
print(f"Error loading image: {str(e)}")
return
else:
# Use a test image URL
test_image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/bert-architects.png"
print(f"Using test image URL: {test_image_url}")
# Download the image
response = requests.get(test_image_url)
image = Image.open(io.BytesIO(response.content))
# Convert image to bytes
img_byte_arr = io.BytesIO()
image.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
# Prepare the request
files = {
'data': ('image.png', img_byte_arr, 'image/png')
}
# Send the request
response = requests.post(server_url, files=files)
print("Sending request to local Gradio server...")
try:
# Check if the request was successful
if response.status_code == 200:
# Parse the response
result = response.json()
if "error" in result:
print(f"Error from server: {result['error']}")
else:
# Extract the embedding from the response
embedding_data = result['data'][0]
embedding_dim = result['data'][1]
print("βœ… Test successful!")
print(f"Embedding dimension: {embedding_dim}")
if embedding_data:
print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
else:
print("No embedding data returned")
else:
print(f"❌ Error: HTTP {response.status_code}")
print(response.text)
except Exception as e:
print(f"❌ Error connecting to server: {str(e)}")
print("Make sure the server is running with 'python app.py'")
if __name__ == "__main__":
# Check if an image path was provided
if len(sys.argv) > 1:
image_path = sys.argv[1]
print(f"Testing with image: {image_path}")
test_local_server(image_path)
else:
print("No image path provided, using test URL")
test_local_server()