File size: 4,332 Bytes
ac44754 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | #!/usr/bin/env python3
"""Smoke-test a deployed GeoText-1652 Inference Endpoint.
Usage:
export HF_TOKEN=hf_...
python test_endpoint.py
Optional:
GEOTEXT_ENDPOINT_URL=https://...endpoints.huggingface.cloud \
python test_endpoint.py --patches --store-patches
"""
from __future__ import annotations
import argparse
import os
import sys
from typing import Any, Dict, Optional
import requests
DEFAULT_ENDPOINT = "https://6ac5162dc8b01bf6eda56cb3.endpoints.huggingface.cloud"
DEFAULT_IMAGE = (
"https://huggingface.co/datasets/geobase/geoai-cogs/resolve/main/"
"geoembeddings-demo/building-detection_sm.tif"
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--endpoint-url",
default=os.environ.get("GEOTEXT_ENDPOINT_URL", DEFAULT_ENDPOINT),
help="Endpoint base URL; defaults to GEOTEXT_ENDPOINT_URL or the deployed endpoint.",
)
parser.add_argument("--image-url", default=DEFAULT_IMAGE, help="Public RGB GeoTIFF/COG URL.")
parser.add_argument(
"--patches", action="store_true", help="Also test tiled patch-embedding output."
)
parser.add_argument(
"--store-patches",
action="store_true",
help="Also test Parquet persistence; endpoint must have HF_TOKEN and HF_BUCKET configured.",
)
return parser.parse_args()
def call(
session: requests.Session,
method: str,
url: str,
payload: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
response = session.request(method, url, json=payload, timeout=300)
try:
response.raise_for_status()
except requests.HTTPError as error:
raise RuntimeError(f"{method} {url} failed: {response.text}") from error
return response.json()
def check(name: str, result: Dict[str, Any], required: set[str]) -> None:
missing = required.difference(result)
if missing:
raise RuntimeError(f"{name} response is missing: {sorted(missing)}\n{result}")
print(f"PASS {name}")
def main() -> None:
args = parse_args()
token = os.environ.get("HF_TOKEN")
if not token:
sys.exit("Set HF_TOKEN to a Hugging Face token with access to the private endpoint.")
base_url = args.endpoint_url.rstrip("/")
session = requests.Session()
session.headers.update(
{"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
)
health = call(session, "GET", f"{base_url}/health")
check("health", health, {"status"})
text = call(
session,
"POST",
f"{base_url}/embed/text",
{"queries": ["a parking lot", "a building"]},
)
check("text embeddings", text, {"embedding_dimension", "text_embeddings"})
assert len(text["text_embeddings"]) == 2
assert len(text["text_embeddings"][0]["embedding"]) == 256
image = call(
session,
"POST",
f"{base_url}/embed/image",
{"image_url": args.image_url, "include_global_embedding": True},
)
check("global image embedding", image, {"image_embedding", "image_metadata"})
assert len(image["image_embedding"]) == 256
ranking = call(
session,
"POST",
f"{base_url}/infer",
{"image_url": args.image_url, "queries": ["a parking lot", "a building"]},
)
check("text-guided ranking", ranking, {"ranked_queries", "image_metadata"})
assert len(ranking["ranked_queries"]) == 2
if args.patches or args.store_patches:
patch_request: Dict[str, Any] = {
"image_url": args.image_url,
"tile_size": 384,
"tile_overlap": 64,
"include_global_embedding": True,
"include_patch_embeddings": args.patches,
"store_patch_embeddings": args.store_patches,
}
patches = call(session, "POST", f"{base_url}/embed/image", patch_request)
check("tiled patch embeddings", patches, {"tiling"})
if args.patches:
check("patch-vector response", patches, {"patch_embeddings"})
assert patches["patch_embeddings"]["patches"]
if args.store_patches:
check("patch-vector storage", patches, {"patch_embeddings_storage"})
print("All requested GeoText endpoint checks passed.")
if __name__ == "__main__":
main()
|