1"""Multi-modal embedder combining different modalities.
2
3This module provides a unified interface for embedding documents across
4multiple modalities (text, image, audio, video) and fusing them into
5combined representations.
6"""
7
8from __future__ import annotations
9
10from typing import cast
11
12import numpy as np
13
14from lexigram.ai.rag.multimodal.embeddings.clip import CLIPEmbedding
15from lexigram.ai.rag.multimodal.types import (
16 AudioDocument,
17 ImageDocument,
18 Modality,
19 MultiModalDocument,
20 MultiModalEmbedding,
21 VideoDocument,
22)
23
24
25class MultiModalEmbedder:
26 """Embedder for multi-modal documents.
27
28 Combines embeddings from different modalities and provides fusion strategies
29 to create unified representations.
30
31 Args:
32 use_clip: Whether to use CLIP for image/text embeddings
33 clip_model: CLIP model name
34 fusion_method: Method for fusing embeddings ("concat", "average", "weighted")
35 weights: Optional weights for weighted fusion
36
37 Example:
38 >>> embedder = MultiModalEmbedder()
39 >>> # Embed multi-modal document
40 >>> embeddings = await embedder.embed(multimodal_doc)
41 >>> print(embeddings.available_modalities)
42 [Modality.TEXT, Modality.IMAGE]
43 """
44
45 def __init__(
46 self,
47 use_clip: bool = True,
48 clip_model: str = "openai/clip-vit-base-patch32",
49 fusion_method: str = "concat",
50 weights: dict[str, float] | None = None,
51 ):
52 """Initialize multi-modal embedder."""
53 self.use_clip = use_clip
54 self.fusion_method = fusion_method
55 self.weights = weights or {
56 "text": 1.0,
57 "image": 1.0,
58 "audio": 1.0,
59 "video": 1.0,
60 }
61
62 # Initialize CLIP if requested
63 self._clip: CLIPEmbedding | None = None
64 if use_clip:
65 self._clip = CLIPEmbedding(model_name=clip_model)
66
67 async def embed_text(self, text: str) -> list[float]:
68 """Embed text content.
69
70 Args:
71 text: Text to embed
72
73 Returns:
74 Text embedding vector
75 """
76 if self._clip:
77 return cast("list[float]", await self._clip.embed_text(text))
78 # Fallback: use simple text embedder
79 msg = "Text embedding requires CLIP or custom embedder"
80 raise NotImplementedError(msg)
81
82 async def embed_image(self, image: ImageDocument) -> list[float]:
83 """Embed image document.
84
85 Args:
86 image: ImageDocument to embed
87
88 Returns:
89 Image embedding vector
90 """
91 if self._clip:
92 return cast("list[float]", await self._clip.embed_image(image))
93 msg = "Image embedding requires CLIP or custom embedder"
94 raise NotImplementedError(msg)
95
96 async def embed_audio(self, audio: AudioDocument) -> list[float]:
97 """Embed audio document.
98
99 For now, uses transcript-based embedding if available.
100
101 Args:
102 audio: AudioDocument to embed
103
104 Returns:
105 Audio embedding vector
106 """
107 # Strategy: Use transcript if available, otherwise placeholder
108 if audio.transcript and self._clip:
109 return cast("list[float]", await self._clip.embed_text(audio.transcript))
110 # Placeholder: could use Wav2Vec or other audio embedder
111 msg = "Direct audio embedding not yet implemented. Enable transcription or provide transcript."
112 raise NotImplementedError(msg)
113
114 async def embed_video(self, video: VideoDocument) -> list[float]:
115 """Embed video document.
116
117 Combines frame embeddings and audio embedding.
118
119 Args:
120 video: VideoDocument to embed
121
122 Returns:
123 Video embedding vector
124 """
125 embeddings = []
126
127 # Embed frames
128 if video.frames and self._clip:
129 frame_embeddings = []
130 for frame in video.frames:
131 frame_emb = await self._clip.embed_image(frame)
132 frame_embeddings.append(frame_emb)
133
134 # Average frame embeddings
135 if frame_embeddings:
136 avg_frame_emb = np.mean(frame_embeddings, axis=0).tolist()
137 embeddings.append(avg_frame_emb)
138
139 # Embed audio/transcript
140 if video.transcript and self._clip:
141 audio_emb = await self._clip.embed_text(video.transcript)
142 embeddings.append(audio_emb)
143 elif video.audio_track:
144 try:
145 audio_emb = await self.embed_audio(video.audio_track)
146 embeddings.append(audio_emb)
147 except NotImplementedError:
148 pass
149
150 # Fuse embeddings
151 if not embeddings:
152 msg = "No embeddings available to fuse"
153 raise ValueError(msg)
154
155 if len(embeddings) == 1:
156 return embeddings[0]
157 return self._fuse_embeddings(embeddings)
158
159 async def embed(self, document: MultiModalDocument) -> MultiModalEmbedding:
160 """Embed a multi-modal document.
161
162 Args:
163 document: MultiModalDocument to embed
164
165 Returns:
166 MultiModalEmbedding with embeddings for each modality
167 """
168 embeddings = {}
169
170 # Embed text
171 if document.text_content:
172 embeddings["text"] = await self.embed_text(document.text_content)
173
174 # Embed images
175 if document.images:
176 image_embeddings = []
177 for image in document.images:
178 img_emb = await self.embed_image(image)
179 image_embeddings.append(img_emb)
180
181 # Average image embeddings
182 embeddings["image"] = np.mean(image_embeddings, axis=0).tolist()
183
184 # Embed audio
185 if document.audio:
186 audio_embeddings = []
187 for audio in document.audio:
188 try:
189 audio_emb = await self.embed_audio(audio)
190 audio_embeddings.append(audio_emb)
191 except NotImplementedError:
192 continue
193
194 if audio_embeddings:
195 embeddings["audio"] = np.mean(audio_embeddings, axis=0).tolist()
196
197 # Embed videos
198 if document.videos:
199 video_embeddings = []
200 for video in document.videos:
201 video_emb = await self.embed_video(video)
202 video_embeddings.append(video_emb)
203
204 embeddings["video"] = np.mean(video_embeddings, axis=0).tolist()
205
206 # Fuse embeddings
207 fused = self._fuse_multimodal_embeddings(embeddings)
208
209 return MultiModalEmbedding(
210 text=embeddings.get("text"),
211 image=embeddings.get("image"),
212 audio=embeddings.get("audio"),
213 video=embeddings.get("video"),
214 fused=fused,
215 fusion_method=self.fusion_method,
216 )
217
218 def _fuse_embeddings(self, embeddings: list[list[float]]) -> list[float]:
219 """Fuse multiple embeddings into one.
220
221 Args:
222 embeddings: List of embedding vectors
223
224 Returns:
225 Fused embedding vector
226 """
227 if self.fusion_method == "concat":
228 # Concatenate all embeddings
229 return np.concatenate(embeddings).tolist()
230
231 if self.fusion_method == "average":
232 # Average all embeddings
233 return np.mean(embeddings, axis=0).tolist()
234
235 if self.fusion_method == "weighted":
236 # Weighted average (requires same dimension)
237 # Not implemented yet
238 msg = "Weighted fusion not implemented"
239 raise NotImplementedError(msg)
240
241 msg = f"Unknown fusion method: {self.fusion_method}"
242 raise ValueError(msg)
243
244 def _fuse_multimodal_embeddings(
245 self,
246 embeddings: dict[str, list[float]],
247 ) -> list[float]:
248 """Fuse embeddings from different modalities.
249
250 Args:
251 embeddings: Dictionary of modality -> embedding
252
253 Returns:
254 Fused embedding vector
255 """
256 if not embeddings:
257 return []
258
259 # Order modalities consistently
260 ordered_modalities = ["text", "image", "audio", "video"]
261 available_embeddings = [
262 embeddings[mod] for mod in ordered_modalities if mod in embeddings
263 ]
264
265 return self._fuse_embeddings(available_embeddings)
266
267 def get_embedding_dimension(self, modality: Modality | None = None) -> int:
268 """Get embedding dimension for a modality.
269
270 Args:
271 modality: Optional modality to get dimension for
272
273 Returns:
274 Embedding dimension
275 """
276 if self._clip:
277 return self._clip.get_embedding_dimension()
278 return 512 # Default CLIP dimension