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