src/clip_embedder.py

from __future__ import annotations

import logging
from functools import lru_cache
from pathlib import Path
from typing import Any

import numpy as np
from PIL import Image

from src.config import DEFAULT_MODEL_NAME, ConfigurationError


logger = logging.getLogger(__name__)


def _detect_device() -> str:
    try:
        import torch

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"


@lru_cache(maxsize=2)
def _load_model(model_name: str, device: str) -> Any:
    try:
        from sentence_transformers import SentenceTransformer
    except ImportError as exc:
        raise ConfigurationError(
            "The sentence-transformers package is not installed. Run `pip install -r requirements.txt`."
        ) from exc

    logger.info("Loading CLIP model %s on %s", model_name, device)
    model = SentenceTransformer(model_name, device=device)
    logger.info("CLIP model loaded")
    return model


class ClipEmbedder:
    def __init__(self, model_name: str = DEFAULT_MODEL_NAME) -> None:
        self.model_name = model_name
        self.device = _detect_device()
        self._model = _load_model(self.model_name, self.device)

    def encode_text(self, text: str) -> list[float]:
        text = (text or "").strip()
        if not text:
            raise ValueError("Text query is empty.")
        logger.info("Encoding text query")
        embedding = self._model.encode(
            [text],
            batch_size=1,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )[0]
        return self._normalize_to_list(embedding)

    def encode_image(self, image: Image.Image | str | Path) -> list[float]:
        logger.info("Encoding image query")
        if isinstance(image, Image.Image):
            pil_image = image.convert("RGB")
        else:
            with Image.open(image) as opened:
                pil_image = opened.convert("RGB")
        embedding = self._model.encode(
            [pil_image],
            batch_size=1,
            convert_to_numpy=True,
            normalize_embeddings=True,
            show_progress_bar=False,
        )[0]
        return self._normalize_to_list(embedding)

    @staticmethod
    def _normalize_to_list(vector: np.ndarray) -> list[float]:
        array = np.asarray(vector, dtype=np.float32).reshape(-1)
        norm = float(np.linalg.norm(array))
        if norm > 0:
            array = array / norm
        return array.astype(float).tolist()