1 files changed, 111 insertions, 0 deletions
diff --git a/embedding/server.py b/embedding/server.py
new file mode 100644
index 0000000..b58f041
--- /dev/null
+++ b/embedding/server.py
@@ -0,0 +1,111 @@
+import base64
+import io
+import os
+import sys
+from typing import Any
+
+import numpy as np
+import torch
+import uvicorn
+from fastapi import FastAPI
+from pydantic import BaseModel
+from PIL import Image
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from embedding.embedder import CompactEmbedder
+
+app = FastAPI()
+model = None
+
+
+@app.on_event("startup")
+def load_model():
+    global model
+    model = CompactEmbedder(local_dir="./models/multi-modal-embed-small")
+
+
+class InlineData(BaseModel):
+    mime_type: str
+    data: str
+
+
+class Part(BaseModel):
+    text: str | None = None
+    inline_data: InlineData | None = None
+
+
+class Content(BaseModel):
+    parts: list[Part]
+
+
+class EmbedContentRequest(BaseModel):
+    model: str = "models/multi-modal-embed-small"
+    content: Content
+    output_dimensionality: int | None = None
+    task_type: str | None = None
+    title: str | None = None
+
+
+class BatchEmbedContentRequest(BaseModel):
+    requests: list[EmbedContentRequest]
+
+
+def _encode_content(content: Content) -> list[float]:
+    vecs: list[Any] = []
+    for part in content.parts:
+        if part.text:
+            vecs.append(model.encode_text(part.text))
+        if part.inline_data:
+            raw = base64.b64decode(part.inline_data.data)
+            mime = part.inline_data.mime_type
+            if mime.startswith("image/"):
+                img = Image.open(io.BytesIO(raw)).convert("RGB")
+                vecs.append(model.encode_image(img))
+            elif mime.startswith("audio/"):
+                audio = np.frombuffer(raw, dtype=np.float32)
+                vecs.append(model.encode_audio(audio))
+    if not vecs:
+        vecs.append(model.encode_text(""))
+    vec = vecs[0] if len(vecs) == 1 else torch.mean(torch.stack(vecs), dim=0)
+    return vec.squeeze(0).tolist()
+
+
+class EmbeddingResponse(BaseModel):
+    embedding: dict[str, list[float]]
+
+
+class EmbedContentResponse(BaseModel):
+    embedding: dict[str, list[float]]
+
+
+class BatchEmbedContentResponse(BaseModel):
+    embeddings: list[dict[str, list[float]]]
+
+
+@app.post("/v1beta/models/{model_id}:embedContent")
+async def embed_content(
+    model_id: str,
+    req: EmbedContentRequest,
+):
+    emb = _encode_content(req.content)
+    if req.output_dimensionality:
+        emb = emb[: req.output_dimensionality]
+    return {"embedding": {"values": emb}}
+
+
+@app.post("/v1beta/models/{model_id}:batchEmbedContents")
+async def batch_embed_contents(
+    model_id: str,
+    req: BatchEmbedContentRequest,
+):
+    embeddings = []
+    for sub_req in req.requests:
+        emb = _encode_content(sub_req.content)
+        if sub_req.output_dimensionality:
+            emb = emb[: sub_req.output_dimensionality]
+        embeddings.append({"values": emb})
+    return {"embeddings": embeddings}
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)