Reference for ultralytics/nn/text_model.py

@abstractmethod
def encode_text(self, texts, dtype):
    """Encode tokenized texts into normalized feature vectors."""
    pass

@abstractmethod
def tokenize(self, texts):
    """Convert input texts to tokens for model processing."""
    pass

@smart_inference_mode()
def encode_image(self, image: Union[Image.Image, torch.Tensor], dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """
    Encode preprocessed images into normalized feature vectors.

    This method processes preprocessed image inputs through the CLIP model to generate feature vectors, which are then
    normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.

    Args:
        image (PIL.Image | torch.Tensor): Preprocessed image input. If a PIL Image is provided, it will be
            converted to a tensor using the model's image preprocessing function.
        dtype (torch.dtype, optional): Data type for output features.

    Returns:
        (torch.Tensor): Normalized image feature vectors with unit length (L2 norm = 1).

    Examples:
        >>> from ultralytics.nn.text_model import CLIP
        >>> from PIL import Image
        >>> clip_model = CLIP("ViT-B/32", device="cuda")
        >>> image = Image.open("path/to/image.jpg")
        >>> image_tensor = clip_model.image_preprocess(image).unsqueeze(0).to("cuda")
        >>> features = clip_model.encode_image(image_tensor)
        >>> features.shape
        torch.Size([1, 512])
    """
    if isinstance(image, Image.Image):
        image = self.image_preprocess(image).unsqueeze(0).to(self.device)
    img_feats = self.model.encode_image(image).to(dtype)
    img_feats = img_feats / img_feats.norm(p=2, dim=-1, keepdim=True)
    return img_feats

@smart_inference_mode()
def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """
    Encode tokenized texts into normalized feature vectors.

    This method processes tokenized text inputs through the CLIP model to generate feature vectors, which are then
    normalized to unit length. These normalized vectors can be used for text-image similarity comparisons.

    Args:
        texts (torch.Tensor): Tokenized text inputs, typically created using the tokenize() method.
        dtype (torch.dtype, optional): Data type for output features.

    Returns:
        (torch.Tensor): Normalized text feature vectors with unit length (L2 norm = 1).

    Examples:
        >>> clip_model = CLIP("ViT-B/32", device="cuda")
        >>> tokens = clip_model.tokenize(["a photo of a cat", "a photo of a dog"])
        >>> features = clip_model.encode_text(tokens)
        >>> features.shape
        torch.Size([2, 512])
    """
    txt_feats = self.model.encode_text(texts).to(dtype)
    txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)
    return txt_feats

def tokenize(self, texts: Union[str, List[str]]) -> torch.Tensor:
    """
    Convert input texts to CLIP tokens.

    Args:
        texts (str | List[str]): Input text or list of texts to tokenize.

    Returns:
        (torch.Tensor): Tokenized text tensor with shape (batch_size, context_length) ready for model processing.

    Examples:
        >>> model = CLIP("ViT-B/32", device="cpu")
        >>> tokens = model.tokenize("a photo of a cat")
        >>> print(tokens.shape)  # torch.Size([1, 77])
    """
    return clip.tokenize(texts).to(self.device)

@smart_inference_mode()
def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """
    Encode tokenized texts into normalized feature vectors.

    Args:
        texts (torch.Tensor): Tokenized text inputs.
        dtype (torch.dtype, optional): Data type for output features.

    Returns:
        (torch.Tensor): Normalized text feature vectors with L2 normalization applied.

    Examples:
        >>> model = MobileCLIP("s0", device="cpu")
        >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
        >>> features = model.encode_text(tokens)
        >>> features.shape
        torch.Size([2, 512])  # Actual dimension depends on model size
    """
    text_features = self.model.encode_text(texts).to(dtype)
    text_features /= text_features.norm(p=2, dim=-1, keepdim=True)
    return text_features

def tokenize(self, texts: List[str]) -> torch.Tensor:
    """
    Convert input texts to MobileCLIP tokens.

    Args:
        texts (List[str]): List of text strings to tokenize.

    Returns:
        (torch.Tensor): Tokenized text inputs with shape (batch_size, sequence_length).

    Examples:
        >>> model = MobileCLIP("s0", "cpu")
        >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
    """
    return self.tokenizer(texts).to(self.device)

@smart_inference_mode()
def encode_text(self, texts: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor:
    """
    Encode tokenized texts into normalized feature vectors.

    Args:
        texts (torch.Tensor): Tokenized text inputs.
        dtype (torch.dtype, optional): Data type for output features.

    Returns:
        (torch.Tensor): Normalized text feature vectors with L2 normalization applied.

    Examples:
        >>> model = MobileCLIPTS(device="cpu")
        >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
        >>> features = model.encode_text(tokens)
        >>> features.shape
        torch.Size([2, 512])  # Actual dimension depends on model size
    """
    # NOTE: no need to do normalization here as it's embedded in the torchscript model
    return self.encoder(texts).to(dtype)

def tokenize(self, texts: List[str]) -> torch.Tensor:
    """
    Convert input texts to MobileCLIP tokens.

    Args:
        texts (List[str]): List of text strings to tokenize.

    Returns:
        (torch.Tensor): Tokenized text inputs with shape (batch_size, sequence_length).

    Examples:
        >>> model = MobileCLIPTS("cpu")
        >>> tokens = model.tokenize(["a photo of a cat", "a photo of a dog"])
    """
    return self.tokenizer(texts).to(self.device)