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
|
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)
|