Skip to content

pixano_inference.models.embedding

Embedding model base class.

A single capability that embeds either an image or text into a shared vector space (CLIP-style), so image and text embeddings are directly comparable (text-to-image search). The I/O types live in :mod:pixano_inference.schemas.embedding and are re-exported here so from pixano_inference.models.embedding import EmbeddingInput keeps working.

EmbeddingModel(config)

Bases: InferenceModel[EmbeddingInput, EmbeddingOutput]

Base class for image/text embedding models (CLIP-style shared space).

predict receives an :class:EmbeddingInput (exactly one of image or text) and returns an :class:EmbeddingOutput (a [num_inputs, dim] vector array). The interface property, when overridden, returns an :class:EmbeddingInterface: the modalities the model embeds and the size of its vectors.

Example 
@register_model("my-embedder")
class MyEmbedder(EmbeddingModel):
    def load_model(self):
        self.model = load_weights(self.config.model_params["path"])

    def predict(self, input: EmbeddingInput) -> EmbeddingOutput:
        vectors = self.model.encode(input.image or input.text)
        return EmbeddingOutput(embeddings=NDArrayFloat.from_numpy(vectors), dim=vectors.shape[-1])
Source code in pixano_inference/models/base.py 
53
54
55
56
57
58
59
def __init__(self, config: ModelDeploymentConfig) -> None:
    """Initialize the model with deployment config.

    Args:
        config: Model deployment configuration.
    """
    self._config = config