1 files changed, 124 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}")