diff options
Diffstat (limited to 'embedding')
| -rw-r--r-- | embedding/embedder.py | 124 | ||||
| -rw-r--r-- | embedding/server.py | 111 |
2 files changed, 235 insertions, 0 deletions
diff --git a/embedding/embedder.py b/embedding/embedder.py new file mode 100644 index 0000000..db353ab --- /dev/null +++ b/embedding/embedder.py @@ -0,0 +1,124 @@ +import os +import torch +import torch.nn as nn +import torch.nn.functional as F +from PIL import Image +from transformers import ( + AutoModel, + BertTokenizer, + SiglipModel, + SiglipProcessor, + WhisperModel, + WhisperFeatureExtractor, +) + + +class CompactEmbedder(nn.Module): + def __init__(self, local_dir="./multi-modal-embed-small"): + local_dir = os.path.abspath(local_dir) + super().__init__() + text_path = os.path.normpath(f"{local_dir}/../base-models/all-MiniLM-L6-v2") + image_path = os.path.normpath( + f"{local_dir}/../base-models/siglip-base-patch16-512" + ) + audio_path = os.path.normpath(f"{local_dir}/../base-models/whisper-tiny") + + self.text_tokenizer = BertTokenizer.from_pretrained( + text_path, local_files_only=True + ) + self.text_encoder = AutoModel.from_pretrained(text_path, local_files_only=True) + + self.image_processor = SiglipProcessor.from_pretrained( + image_path, local_files_only=True + ) + self.image_encoder = SiglipModel.from_pretrained( + image_path, local_files_only=True + ).vision_model + self.audio_processor = WhisperFeatureExtractor.from_pretrained( + audio_path, local_files_only=True + ) + self.image_proj = nn.Linear(768, 384) + self.audio_processor = WhisperFeatureExtractor.from_pretrained( + audio_path, local_files_only=True + ) + self.audio_encoder = WhisperModel.from_pretrained( + audio_path, local_files_only=True + ).encoder + + checkpoint_path = f"{local_dir}/model.pt" + state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) + + self.text_encoder.load_state_dict( + { + k.replace("text_encoder.encoder.", ""): v + for k, v in state_dict.items() + if k.startswith("text_encoder.encoder.") + } + ) + + self.image_encoder.load_state_dict( + { + k.replace("image_encoder.vision_encoder.", ""): v + for k, v in state_dict.items() + if k.startswith("image_encoder.vision_encoder.") + } + ) + + self.image_proj.load_state_dict( + { + k.replace("image_encoder.projection.", ""): v + for k, v in state_dict.items() + if k.startswith("image_encoder.projection.") + } + ) + + self.audio_encoder.load_state_dict( + { + k.replace("audio_encoder.encoder.", ""): v + for k, v in state_dict.items() + if k.startswith("audio_encoder.encoder.") + } + ) + + self.eval() + + @torch.no_grad() + def encode_text(self, texts): + if isinstance(texts, str): + texts = [texts] + inputs = self.text_tokenizer( + texts, padding=True, truncation=True, return_tensors="pt" + ) + inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()} + outputs = self.text_encoder(**inputs) + embeddings = outputs.last_hidden_state.mean(dim=1) + return F.normalize(embeddings, p=2, dim=-1) + + @torch.no_grad() + def encode_image(self, images): + if isinstance(images, Image.Image): + images = [images] + inputs = self.image_processor(images=images, return_tensors="pt") + inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()} + outputs = self.image_encoder(**inputs) + embeddings = self.image_proj(outputs.pooler_output) + return F.normalize(embeddings, p=2, dim=-1) + + @torch.no_grad() + def encode_audio(self, waveform): + if isinstance(waveform, torch.Tensor): + waveform = waveform.squeeze().numpy() + inputs = self.audio_processor( + waveform, sampling_rate=16000, return_tensors="pt" + ) + inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()} + outputs = self.audio_encoder(**inputs) + embeddings = outputs.last_hidden_state.mean(dim=1) + return F.normalize(embeddings, p=2, dim=-1) + + +if __name__ == "__main__": + model = CompactEmbedder(local_dir="./models/multi-modal-embed-small") + + text_vec = model.encode_text("example query") + print(f"Text shape: {text_vec.shape}") 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) |