« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""CLIP embeddings for cross-modal image-text retrieval.
3This module provides CLIP (Contrastive Language-Image Pre-training) embeddings
4that create aligned vector spaces for images and text, enabling:
5- Text-to-image search
6- Image-to-text search
7- Image similarity search
8- Zero-shot image classification
9"""
11from __future__ import annotations
13from pathlib import Path
14import sys
15from typing import TYPE_CHECKING, Any, cast
17if TYPE_CHECKING:
18 from PIL.Image import Image as PILImage
20 PILImageType = PILImage
21else:
22 PILImageType: Any = Any
24# Runtime optional dependency enforcement: imports happen lazily inside the
25# CLIPEmbedding constructor to avoid import-time failures in environments
26# where heavy ML packages are not installed. Use helper to give actionable
27# error messages when required packages are missing.
29from lexigram.ai.rag.deps import MissingOptionalDependencyError, ensure_packages
32def __getattr__(name: str) -> Any:
33 """Allow module-level attribute assignments for lazy optional imports.
35 This satisfies mypy while still allowing runtime lazy loading of CLIP
36 dependencies. The actual attributes (CLIPModel, CLIPProcessor, torch,
37 Image) are set at runtime in the constructor.
38 """
39 if name in {"CLIPModel", "CLIPProcessor", "torch", "Image"}:
40 try:
41 return sys.modules[__name__].__dict__[name]
42 except KeyError:
43 pass
44 msg = f"module {__name__!r} has no attribute {name!r}"
45 raise AttributeError(msg)
48# Module-level placeholders populated at runtime when the packages are available
49torch: Any = None
50Image: Any = None
51CLIPModel: Any | None = None
52CLIPProcessor: Any | None = None
54# Backwards-compat shim for existing tests that patch this variable
55CLIP_AVAILABLE = False
57from lexigram.ai.rag.exceptions import CLIPEmbeddingError
58from lexigram.ai.rag.multimodal.types import ImageDocument
61class CLIPEmbedding:
62 """CLIP embeddings for images and text.
64 Uses OpenAI's CLIP model to create aligned embeddings for images and text
65 in the same vector space, enabling cross-modal retrieval.
67 Args:
68 model_name: CLIP model name from HuggingFace
69 device: Device to run model on ("cpu", "cuda", "mps")
70 batch_size: Batch size for processing multiple items
72 Example:
73 >>> embedder = CLIPEmbedding()
74 >>> # Embed text
75 >>> text_emb = await embedder.embed_text("a photo of a cat")
76 >>> # Embed image
77 >>> img_emb = await embedder.embed_image(image_doc)
78 >>> # Compute similarity
79 >>> similarity = cosine_similarity(text_emb, img_emb)
80 """
82 def __init__(
83 self,
84 model_name: str = "openai/clip-vit-base-patch32",
85 device: str | None = None,
86 batch_size: int = 32,
87 ):
88 """Initialize CLIP embedder."""
89 # If test suite or older callers set CLIP_AVAILABLE to True, honor that
90 # and avoid importing the heavy runtime packages (they will be patched
91 # in tests). Otherwise, ensure the optional packages are present and
92 # perform lazy imports.
93 mod = sys.modules[__name__]
94 if not CLIP_AVAILABLE:
95 try:
96 _ = ensure_packages(
97 ["transformers", "torch", "PIL"],
98 hint="pip install 'lexigram-ai[llm]' or 'lexigram-ai[all]'",
99 )
100 except MissingOptionalDependencyError as e:
101 msg = (
102 "CLIP embeddings require transformers, torch, and pillow. "
103 "Install with: pip install 'lexigram-ai[llm]' or 'lexigram-ai[all]'."
104 )
105 raise ImportError(msg) from e
107 # Perform actual imports now that availability is asserted.
108 # Wrap imports and re-raise a consistent ImportError so tests and
109 # callers see a helpful message rather than package internals.
110 try:
111 from PIL import Image as _Image
112 import torch as _torch
113 from transformers import (
114 CLIPModel as _CLIPModel,
115 )
116 from transformers import CLIPProcessor as _CLIPProcessor
117 except Exception as e:
118 msg = (
119 "CLIP embeddings require transformers, torch, and pillow. "
120 "Install with: pip install 'lexigram-ai[llm]' or 'lexigram-ai[all]'."
121 )
122 raise ImportError(msg) from e
124 mod.CLIPModel = _CLIPModel # type: ignore[attr-defined]
125 mod.CLIPProcessor = _CLIPProcessor # type: ignore[attr-defined]
126 # Also update the local module-level bindings so that subsequent
127 # references (e.g. `torch.cuda`) work correctly within this process.
128 torch = _torch
129 Image = _Image
130 else:
131 # CLIP_AVAILABLE True -> assume tests or caller patched module-level
132 # symbols (CLIPModel/CLIPProcessor/torch/Image). If one is missing,
133 # let normal attribute access/patching handle the failure in tests.
134 torch = sys.modules.get("lexigram.ai.rag.multimodal.embeddings.clip").torch # type: ignore[union-attr]
136 self.model_name = model_name
137 self.batch_size = batch_size
139 # Auto-detect device
140 if device is None:
141 cuda_mod = getattr(torch, "cuda", None)
142 mps_mod = getattr(getattr(torch, "backends", None), "mps", None)
143 if cuda_mod and getattr(cuda_mod, "is_available", lambda: False)():
144 device = "cuda"
145 elif mps_mod and getattr(mps_mod, "is_available", lambda: False)():
146 device = "mps"
147 else:
148 device = "cpu"
149 self.device = device
151 # Load model and processor
152 self._model: Any | None = None
153 self._processor: Any | None = None
154 self._load_model()
156 def _load_model(self) -> Any:
157 """Load CLIP model and processor."""
158 try:
159 self._processor = CLIPProcessor.from_pretrained(self.model_name) # type: ignore[union-attr]
160 self._model = CLIPModel.from_pretrained(self.model_name) # type: ignore[union-attr]
161 self._model.to(self.device)
162 self._model.eval()
163 except Exception as e:
164 msg = f"Failed to load CLIP model: {e}"
165 raise CLIPEmbeddingError(msg) from e
167 async def embed_text(
168 self,
169 text: str | list[str],
170 ) -> list[float] | list[list[float]]:
171 """Embed text using CLIP text encoder.
173 Args:
174 text: Single text string or list of texts
176 Returns:
177 Embedding vector(s) of dimension 512
178 """
179 if not self._model or not self._processor:
180 msg = "CLIP model not loaded or processor missing"
181 raise CLIPEmbeddingError(msg)
183 # Ensure list
184 is_single = isinstance(text, str)
185 texts = [text] if is_single else text
187 try:
188 # Process text
189 inputs = self._processor(
190 text=texts,
191 return_tensors="pt",
192 padding=True,
193 truncation=True,
194 )
195 inputs = {k: v.to(self.device) for k, v in inputs.items()}
197 # Get embeddings
198 with torch.no_grad():
199 text_features = self._model.get_text_features(**inputs)
201 # Normalize embeddings
202 text_features = text_features / text_features.norm(dim=-1, keepdim=True)
204 # Convert to list
205 embeddings = text_features.cpu().numpy().tolist()
207 # Return single embedding or list
208 return embeddings[0] if is_single else embeddings
210 except Exception as e:
211 msg = f"Failed to embed text: {e}"
212 raise CLIPEmbeddingError(msg) from e
214 async def embed_image(
215 self,
216 image: (
217 ImageDocument
218 | PILImageType
219 | str
220 | Path
221 | list[ImageDocument | PILImageType | str | Path]
222 ),
223 ) -> list[float] | list[list[float]]:
224 """Embed image using CLIP image encoder.
226 Args:
227 image: ImageDocument, PIL Image, file path, or list of any
229 Returns:
230 Embedding vector(s) of dimension 512
231 """
232 if not self._model or not self._processor:
233 msg = "CLIP model not loaded or processor missing"
234 raise CLIPEmbeddingError(msg)
236 # Normalize to a list for processing
237 images: list[ImageDocument | PILImageType | str | Path]
238 if isinstance(image, list):
239 images = image
240 is_single = False
241 else:
242 images = [image]
243 is_single = True
245 try:
246 # Load images
247 pil_images = []
248 for img in images:
249 pil_img = self._load_image(img)
250 pil_images.append(pil_img)
252 # Process images
253 inputs = self._processor(
254 images=pil_images,
255 return_tensors="pt",
256 )
257 inputs = {k: v.to(self.device) for k, v in inputs.items()}
259 # Get embeddings
260 with torch.no_grad():
261 image_features = self._model.get_image_features(**inputs)
263 # Normalize embeddings
264 image_features = image_features / image_features.norm(
265 dim=-1,
266 keepdim=True,
267 )
269 # Convert to list
270 embeddings = image_features.cpu().numpy().tolist()
272 # Return single embedding or list
273 return embeddings[0] if is_single else embeddings
275 except Exception as e:
276 msg = f"Failed to embed image: {e}"
277 raise CLIPEmbeddingError(msg) from e
279 async def embed_batch(
280 self,
281 texts: list[str] | None = None,
282 images: list[ImageDocument | PILImage | str | Path] | None = None,
283 ) -> dict:
284 """Embed batches of texts and/or images.
286 Args:
287 texts: Optional list of texts to embed
288 images: Optional list of images to embed
290 Returns:
291 Dictionary with "text_embeddings" and/or "image_embeddings"
292 """
293 result = {}
295 if texts:
296 result["text_embeddings"] = await self.embed_text(texts)
298 if images:
299 result["image_embeddings"] = await self.embed_image(images)
301 return result
303 def _load_image(
304 self,
305 image: ImageDocument | PILImageType | str | Path,
306 ) -> PILImageType:
307 """Load image from various sources.
309 Args:
310 image: ImageDocument, PIL Image, or file path
312 Returns:
313 PIL Image
314 """
315 # Check if it's a PIL Image by checking module and class name
316 if (
317 hasattr(image, "__class__")
318 and image.__class__.__module__.startswith("PIL")
319 and image.__class__.__name__ == "Image"
320 ):
321 return cast("PILImageType", image)
323 if isinstance(image, ImageDocument):
324 # Load from ImageDocument
325 if isinstance(image.content, bytes):
326 import io
328 return cast("PILImageType", Image.open(io.BytesIO(image.content)))
329 return cast("PILImageType", Image.open(image.content))
331 if isinstance(image, (str, Path)):
332 # Load from file path
333 return cast("PILImageType", Image.open(image))
335 msg = f"Unsupported image type: {type(image)}"
336 raise CLIPEmbeddingError(msg)
338 def get_embedding_dimension(self) -> int:
339 """Get the dimension of CLIP embeddings.
341 Returns:
342 Embedding dimension (512 for base model)
343 """
344 return self._model.config.projection_dim if self._model else 512
346 def compute_similarity(
347 self,
348 embedding1: list[float],
349 embedding2: list[float],
350 ) -> float:
351 """Compute cosine similarity between two embeddings.
353 Args:
354 embedding1: First embedding vector
355 embedding2: Second embedding vector
357 Returns:
358 Cosine similarity score between -1 and 1
359 """
360 import numpy as np
362 # Convert to numpy arrays
363 emb1 = np.array(embedding1)
364 emb2 = np.array(embedding2)
366 # Compute cosine similarity
367 similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
369 return float(similarity)