src/models.py

from pydantic import BaseModel, validator
from typing import List
from PIL.Image import Image

from .utils import download_image_as_pil


BATCH_SIZE_TEXT: int = 128
BATCH_SIZE_IMAGE: int = 64


class TextRequest(BaseModel):
    texts: List[str]

    @validator("texts")
    def validate_texts_batch_size(cls, v):
        if len(v) > BATCH_SIZE_TEXT:
            raise ValueError(f"Maximum batch size for texts is {BATCH_SIZE_TEXT}")
        if len(v) == 0:
            raise ValueError("At least one text is required")
        return v


class ImageRequest(BaseModel):
    urls: List[str]

    @validator("urls")
    def validate_images_batch_size(cls, v):
        if len(v) > BATCH_SIZE_IMAGE:
            raise ValueError(f"Maximum batch size for images is {BATCH_SIZE_IMAGE}")
        if len(v) == 0:
            raise ValueError("At least one image URL is required")
        return v

    def download(self) -> List[Image]:
        return [download_image_as_pil(url) for url in self.urls]


class Response(BaseModel):
    embeddings: List[List[float]]