File size: 5,952 Bytes
9037c05 | 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 | #!/usr/bin/env python3
"""OpenAI-compatible embeddings server for EmbeddingGemma 2 on an LLM-8850 (AXCL) card.
python3 eg2_server.py --host 0.0.0.0 --port 8010
POST /v1/embeddings
{"input": "text" | ["text", ...] | [{"text": "...", "image": "<base64 or data URL>", "audio": "<base64 wav>"}],
"dimensions": 768|512|256|128, "encoding_format": "float"|"base64",
"input_type": "query"|"document"|"sts"|"classification"|"clustering"|"code"} (or "prompt_name")
GET /v1/models, /health
Object inputs may hold several images/audios as lists; placeholders <|image|>/<|audio|> in "text"
mark where they go (otherwise they are appended). Local file paths in "image"/"audio" are only
accepted with --allow-local-files. One request runs at a time (the card is shared).
"""
import argparse
import base64
import io
import json
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import numpy as np
from eg2_host import Eg2, load_audio
MODEL_ID = "embeddinggemma-2-axcl"
def decode_media(value, kind, allow_files):
from PIL import Image
if isinstance(value, str) and value.startswith("data:"):
value = value.split(",", 1)[1]
if isinstance(value, str) and allow_files and Path(value).exists():
return Image.open(value) if kind == "image" else load_audio(value)
raw = base64.b64decode(value)
if kind == "image":
return Image.open(io.BytesIO(raw))
import soundfile as sf
wav, sr = sf.read(io.BytesIO(raw), dtype="float32", always_2d=True)
wav = wav.mean(1)
if sr != 16000:
from math import gcd
from scipy.signal import resample_poly
g = gcd(sr, 16000)
wav = resample_poly(wav, 16000 // g, sr // g).astype(np.float32)
return wav
def make_handler(eg, lock, allow_files):
class Handler(BaseHTTPRequestHandler):
def _send(self, code, obj):
body = json.dumps(obj, ensure_ascii=False).encode()
self.send_response(code)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def do_GET(self):
if self.path.rstrip("/") in ("/v1/models", "/models"):
self._send(200, {"object": "list", "data": [{"id": MODEL_ID, "object": "model", "owned_by": "local"}]})
elif self.path.rstrip("/") == "/health":
self._send(200, {"status": "ok"})
else:
self._send(404, {"error": {"message": "not found"}})
def do_POST(self):
if self.path.rstrip("/") not in ("/v1/embeddings", "/embeddings"):
return self._send(404, {"error": {"message": "not found"}})
try:
req = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0))) or b"{}")
inputs = req.get("input")
if isinstance(inputs, (str, dict)):
inputs = [inputs]
if not inputs:
raise ValueError("input is required")
dim = int(req.get("dimensions") or 768)
prompt = req.get("input_type") or req.get("prompt_name")
out, tokens = [], 0
with lock:
for i, item in enumerate(inputs):
if isinstance(item, str):
v = eg.encode(text=item, prompt=prompt, dim=dim)
else:
imgs = item.get("image") or []
auds = item.get("audio") or []
imgs = [imgs] if isinstance(imgs, str) else imgs
auds = [auds] if isinstance(auds, str) else auds
v = eg.encode(text=item.get("text", ""),
images=[decode_media(x, "image", allow_files) for x in imgs],
audios=[decode_media(x, "audio", allow_files) for x in auds],
prompt=prompt if item.get("text") else None, dim=dim)
tokens += eg.last_tokens
emb = (base64.b64encode(v.astype("<f4").tobytes()).decode()
if req.get("encoding_format") == "base64" else v.tolist())
out.append({"object": "embedding", "index": i, "embedding": emb})
self._send(200, {"object": "list", "data": out, "model": MODEL_ID,
"usage": {"prompt_tokens": tokens, "total_tokens": tokens}})
except Exception as e: # report bad requests as 400 with the reason
self._send(400, {"error": {"message": f"{type(e).__name__}: {e}", "type": "invalid_request_error"}})
def log_message(self, fmt, *args):
print(time.strftime("%H:%M:%S"), self.address_string(), fmt % args)
return Handler
def main():
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
ap.add_argument("--host", default="127.0.0.1")
ap.add_argument("--port", type=int, default=8010)
ap.add_argument("--assets", default=str(Path(__file__).resolve().parent / "assets"))
ap.add_argument("--text-only", action="store_true", help="do not load the vision/audio encoders")
ap.add_argument("--allow-local-files", action="store_true")
args = ap.parse_args()
eg = Eg2(args.assets, load=("text",) if args.text_only else ("text", "vision", "audio"))
server = ThreadingHTTPServer((args.host, args.port), make_handler(eg, threading.Lock(), args.allow_local_files))
print(f"serving {MODEL_ID} on http://{args.host}:{args.port}/v1/embeddings")
try:
server.serve_forever()
except KeyboardInterrupt:
pass
finally:
eg.close()
if __name__ == "__main__":
main()
|