File size: 5,592 Bytes
83d47de | 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 135 136 137 138 139 140 141 | """
Custom inference handler for Hugging Face Inference Endpoints.
Serves the Marqo GCL ViT-L-14 open_clip checkpoint as a pure encoder:
given text and/or image(s), it returns L2-normalized embeddings, exactly
matching the notebook's encode_text / encode_image + normalize steps.
Request (JSON):
{ "inputs": { "text": ["a red dress", ...], # optional
"image": ["<base64-encoded-image>", ...] } } # optional
Response (JSON):
{ "text_embeddings": [[... 768 floats ...], ...], # if text sent
"image_embeddings": [[... 768 floats ...], ...] } # if image sent
Fusion (alpha*img + (1-alpha)*txt) is intentionally NOT done here — that
happens in the n8n Code node, so this endpoint stays a clean encoder.
"""
import os
import io
import glob
import base64
import torch
import open_clip
from PIL import Image
class EndpointHandler:
def __init__(self, path: str = ""):
# `path` is the local directory where the repo (including the .pt) is mounted.
self.device = "cuda" if torch.cuda.is_available() else "cpu"
# Architecture of the checkpoint. ViT-L-14 for the Marqo GCL model.
arch = os.environ.get("OPEN_CLIP_ARCH", "ViT-L-14")
# Build the architecture WITHOUT pretrained weights, then load the GCL checkpoint.
self.model, _, self.preprocess = open_clip.create_model_and_transforms(
arch, pretrained=None
)
self.tokenizer = open_clip.get_tokenizer(arch)
ckpt_path = self._find_checkpoint(path)
self._load_weights(ckpt_path)
self.model = self.model.to(self.device).eval()
# ---- locate the checkpoint file inside the repo ----
def _find_checkpoint(self, path: str) -> str:
# Allow an explicit override via env var (filename or absolute path).
env = os.environ.get("CHECKPOINT_FILE")
if env:
return env if os.path.isabs(env) else os.path.join(path, env)
candidates = (
glob.glob(os.path.join(path, "*.pt"))
+ glob.glob(os.path.join(path, "*.bin"))
+ glob.glob(os.path.join(path, "*.safetensors"))
)
if not candidates:
raise FileNotFoundError(
f"No .pt/.bin/.safetensors checkpoint found in '{path}'. "
"Set CHECKPOINT_FILE env var to the filename."
)
# If several exist, prefer the largest (the model weights, not a config shard).
return max(candidates, key=os.path.getsize)
# ---- load weights, tolerating the 'full_states' training-checkpoint layout ----
def _load_weights(self, ckpt_path: str) -> None:
# The Marqo file is named "...full_states.pt": a full training state that
# wraps the model weights under a 'state_dict' key and may prefix keys with
# 'module.' (from DDP training). open_clip.load_checkpoint normalizes both.
try:
open_clip.load_checkpoint(self.model, ckpt_path, strict=False)
return
except Exception as e:
print(f"[handler] open_clip.load_checkpoint fallback triggered: {e}")
# Manual fallback. weights_only=False is required because a full training
# state contains non-tensor objects (optimizer state, epoch, etc.).
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
if isinstance(ckpt, dict):
state_dict = ckpt.get("state_dict", ckpt.get("model", ckpt))
else:
state_dict = ckpt
# Strip a leading 'module.' from DDP-saved keys.
state_dict = {
(k[len("module."):] if k.startswith("module.") else k): v
for k, v in state_dict.items()
}
missing, unexpected = self.model.load_state_dict(state_dict, strict=False)
if missing:
print(f"[handler] missing keys (first 5): {list(missing)[:5]}")
if unexpected:
print(f"[handler] unexpected keys (first 5): {list(unexpected)[:5]}")
@torch.no_grad()
def __call__(self, data: dict) -> dict:
# HF passes {"inputs": ..., "parameters": ...}; be lenient about shape.
inputs = data.get("inputs", data) or {}
if not isinstance(inputs, dict):
return {"error": "Send an object under 'inputs' with 'text' and/or 'image'."}
out = {}
# ---- text embeddings ----
texts = inputs.get("text")
if texts:
if isinstance(texts, str):
texts = [texts]
tokens = self.tokenizer(texts).to(self.device)
feats = self.model.encode_text(tokens)
feats = feats / feats.norm(dim=-1, keepdim=True)
out["text_embeddings"] = feats.cpu().tolist()
# ---- image embeddings (base64-encoded bytes) ----
images = inputs.get("image")
if images:
if isinstance(images, str):
images = [images]
tensors = []
for b64 in images:
raw = base64.b64decode(b64)
img = Image.open(io.BytesIO(raw)).convert("RGB")
tensors.append(self.preprocess(img))
batch = torch.stack(tensors).to(self.device)
feats = self.model.encode_image(batch)
feats = feats / feats.norm(dim=-1, keepdim=True)
out["image_embeddings"] = feats.cpu().tolist()
if not out:
return {
"error": "Provide 'text' and/or 'image' (base64) under 'inputs'.",
"example": {"inputs": {"text": ["a red dress"]}},
}
return out
|