| | from __future__ import annotations |
| |
|
| | import numpy as np |
| | import torch |
| | from transformers import CLIPModel, CLIPProcessor |
| |
|
| |
|
| | class TextEmbedder: |
| | """ |
| | CLIP projected text features (512-d). |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | model_name: str = "openai/clip-vit-base-patch32", |
| | device: str = "cpu", |
| | ): |
| | self.device = device |
| | self.processor = CLIPProcessor.from_pretrained(model_name) |
| | self.model = CLIPModel.from_pretrained(model_name).to(self.device) |
| | self.model.eval() |
| |
|
| | @torch.no_grad() |
| | def embed(self, text: str) -> np.ndarray: |
| | inputs = self.processor( |
| | text=[text], |
| | return_tensors="pt", |
| | padding=True, |
| | truncation=True, |
| | ).to(self.device) |
| | feats = self.model.get_text_features(**inputs) |
| | |
| | target_dim = getattr(self.model.config, "projection_dim", 512) |
| | if not isinstance(feats, torch.Tensor): |
| | pooled = feats.pooler_output |
| | if pooled.shape[-1] != target_dim: |
| | proj = getattr(self.model, "text_projection", None) |
| | if proj is not None: |
| | pooled = proj(pooled) |
| | feats = pooled |
| | if feats.dim() == 3: |
| | pooled = feats[:, 0, :] |
| | if pooled.shape[-1] != target_dim: |
| | proj = getattr(self.model, "text_projection", None) |
| | if proj is not None: |
| | pooled = proj(pooled) |
| | feats = pooled |
| | if feats.dim() == 2: |
| | feats = feats[0] |
| | return feats.cpu().numpy().astype("float32") |
| |
|
