LabelProcessor.py

import os
from transformers import CLIPProcessor, CLIPModel
import torch

class LabelProcessor:
    def __init__(self):
        """Load CLIP model and processor for computing text embeddings.

        The model path can be provided via the `CLIP_MODEL` environment variable
        or will default to the parent `models/clip-vit-large-patch14` directory.
        """
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        model_path = os.path.normpath(
            os.path.join(os.path.dirname(__file__), "..", "models", "clip-vit-large-patch14")
        )
        self.model = CLIPModel.from_pretrained(model_path).to(self.device)
        self.processor = CLIPProcessor.from_pretrained(model_path)

    def compute_embedding(self, label):
        """Compute a single text embedding for a label string.

        Args:
            label (str): Text label to embed.

        Returns:
            numpy.ndarray: 1-D embedding vector on CPU.
        """
        inputs = self.processor(text=label, return_tensors="pt", padding=True).to(self.device)
        with torch.no_grad():
            outputs = self.model.get_text_features(**inputs)
            pooled = outputs.pooler_output
            embedding = self.model.text_projection(pooled)

        final_embedding = embedding.squeeze(0).cpu().numpy()

        del inputs, outputs, pooled, embedding
        torch.cuda.empty_cache()

        return final_embedding

    def compute_embeddings(self, label):
        """Compute embeddings for a comma-separated label string.

        Splits `label` on commas and computes an embedding for each part.

        Args:
            label (str): Comma-separated label string.

        Returns:
            list[numpy.ndarray]: List of embeddings for each sub-label.
        """
        labels = [part.strip() for part in label.split(",") if part.strip()]

        embeddings = []
        for item in labels:
            embeddings.append(self.compute_embedding(item))

        return embeddings

    def compute_similarity(self, embedding1, embedding2):
        """Compute cosine similarity between two embeddings.

        Args:
            embedding1, embedding2 (array-like): Vectors to compare.

        Returns:
            float: Cosine similarity score between -1 and 1.
        """
        embedding1 = torch.as_tensor(embedding1)
        embedding2 = torch.as_tensor(embedding2)

        similarity = torch.nn.functional.cosine_similarity(
            embedding1,
            embedding2, 
            dim=0
        )
        return similarity.mean().item()