from abc import ABC, abstractmethod

from dataclasses import dataclass

from typing import (

TYPE_CHECKING,

Any,

Callable,

List,

Optional,

Protocol,

runtime_checkable,

)

import pandas as pd

if TYPE_CHECKING:

import numpy as np

from feast.repo_config import EmbeddingModelConfig

@runtime_checkable

class EmbeddingProvider(Protocol):

"""Protocol for query-time embedding providers.

Implement this to plug in any embedding backend (OpenAI, Cohere,

SentenceTransformers, a local model, etc.) for use in

``openai_search`` and the OpenAI-compatible

feature server endpoints.

Example::

class MyEmbeddingProvider:

def embed(self, texts: List[str]) -> List[List[float]]:

return my_model.encode(texts)

async def aembed(self, texts: List[str]) -> List[List[float]]:

return await my_model.aencode(texts)

"""

def embed(self, texts: List[str]) -> List[List[float]]:

"""Return one embedding vector per input text."""

...

async def aembed(self, texts: List[str]) -> List[List[float]]:

"""Async variant of :meth:`embed`."""

...

class SentenceTransformersEmbeddingProvider:

"""EmbeddingProvider backed by Sentence Transformers (local inference).

Runs entirely on-device — no external API key required. Ideal for

air-gapped deployments, cost-sensitive workloads, or environments where

outbound API calls are restricted.

Requires the ``sentence-transformers`` package::

pip install sentence-transformers

Configuration in ``feature_store.yaml``::

embedding_model:

provider: sentence_transformers

model: all-MiniLM-L6-v2 # any HuggingFace sentence-transformers model

The model is loaded lazily on the first call to :meth:`embed` or

:meth:`aembed` so that import cost is paid only when the provider is

actually used.

"""

def __init__(self, model: str = "all-MiniLM-L6-v2"):

self.model_name = model

self._model = None

@classmethod

def from_config(

cls, config: "EmbeddingModelConfig"

) -> "SentenceTransformersEmbeddingProvider":

"""Build a provider from a :class:`~feast.repo_config.EmbeddingModelConfig`."""

return cls(model=config.model)

def _load_model(self):

if self._model is None:

try:

from sentence_transformers import SentenceTransformer

except ImportError:

raise ImportError(

"sentence-transformers is required for the "

"'sentence_transformers' embedding provider. "

"Install with: pip install sentence-transformers"

)

self._model = SentenceTransformer(self.model_name)

return self._model

def embed(self, texts: List[str]) -> List[List[float]]:

model = self._load_model()

embeddings = model.encode(texts, convert_to_numpy=True)

return [emb.tolist() for emb in embeddings]

async def aembed(self, texts: List[str]) -> List[List[float]]:

import asyncio

loop = asyncio.get_event_loop()

return await loop.run_in_executor(None, self.embed, texts)

def get_embedding_provider(

config: "EmbeddingModelConfig",

) -> "EmbeddingProvider":

"""Factory that returns the appropriate :class:`EmbeddingProvider` for *config*.

Dispatches on ``config.provider``:

* ``"sentence_transformers"`` (default) → :class:`SentenceTransformersEmbeddingProvider`

Raises:

ValueError: If ``config.provider`` is not a recognised value.

"""

provider = (config.provider or "sentence_transformers").lower()

if provider == "sentence_transformers":

return SentenceTransformersEmbeddingProvider.from_config(config)

raise ValueError(

f"Unknown embedding provider: '{config.provider}'. "

"Supported values: 'sentence_transformers'."

)

@dataclass

class EmbeddingConfig:

batch_size: int = 64

show_progress: bool = True

class BaseEmbedder(ABC):

"""

Abstract base class for embedding generation.

Supports multiple modalities via routing.

Users can register custom modality handlers.

"""

def __init__(self, config: Optional[EmbeddingConfig] = None):

self.config = config or EmbeddingConfig()

# Registry: modality -> embedding function

self._modality_handlers: dict[str, Callable[[List[Any]], "np.ndarray"]] = {}

# Register default modalities (subclass can override)

self._register_default_modalities()

def _register_default_modalities(self) -> None:

"""Override in subclass to register default modality handlers."""

pass

def register_modality(

self,

modality: str,

handler: Callable[[List[Any]], "np.ndarray"],

) -> None:

"""

Register a handler for a modality.

Args:

modality: Name of modality ("text", "image", "video", etc.)

handler: Function that takes list of inputs and returns embeddings.

"""

self._modality_handlers[modality] = handler

@property

def supported_modalities(self) -> List[str]:

"""Return list of supported modalities."""

return list(self._modality_handlers.keys())

def get_embedding_dim(self, modality: str) -> Optional[int]:

"""

Return the embedding dimension for a given modality.

Subclasses should override this to return the actual dimension

so that auto-generated FeatureView schemas use the correct vector_length.

Args:

modality: The modality to query (e.g. "text", "image").

Returns:

The embedding dimension, or None if unknown.

"""

return None

@abstractmethod

def embed(self, inputs: List[Any], modality: str) -> "np.ndarray":

"""

Generate embeddings for inputs of a given modality.

Args:

inputs: List of inputs.

modality: Type of content ("text", "image", "video", etc.)

Returns:

numpy array of shape (len(inputs), embedding_dim)

"""

pass

def embed_dataframe(

self,

df: pd.DataFrame,

column_mapping: dict[str, tuple[str, str]],

) -> pd.DataFrame:

"""

Add embeddings for multiple columns with modality routing.

Args:

df: Input DataFrame.

column_mapping: Dict mapping source_column -> (modality, output_column).

Example: {

"text": ("text", "text_embedding"),

"image_path": ("image", "image_embedding"),

"video_path": ("video", "video_embedding"),

}

"""

df = df.copy()

for source_column, (modality, output_column) in column_mapping.items():

inputs = df[source_column].tolist()

embeddings = self.embed(inputs, modality)

df[output_column] = pd.Series(

[emb.tolist() for emb in embeddings], dtype=object, index=df.index

)

return df

class MultiModalEmbedder(BaseEmbedder):

"""

Multi-modal embedder with built-in support for common modalities.

Supports: text, image, video (extensible)

"""

def __init__(

self,

text_model: str = "all-MiniLM-L6-v2",

image_model: str = "openai/clip-vit-base-patch32",

config: Optional[EmbeddingConfig] = None,

):

self.text_model_name = text_model

self.image_model_name = image_model

# Lazy-loaded models

self._text_model = None

self._image_model = None

self._image_processor = None

super().__init__(config)

def _register_default_modalities(self) -> None:

"""Register built-in modality handlers."""

self.register_modality("text", self._embed_text)

self.register_modality("image", self._embed_image)

# Future: add more as needed

# self.register_modality("video", self._embed_video)

# self.register_modality("audio", self._embed_audio)

def embed(self, inputs: List[Any], modality: str) -> "np.ndarray":

"""Route to appropriate handler based on modality."""

if modality not in self._modality_handlers:

raise ValueError(

f"Unsupported modality: '{modality}'. "

f"Supported: {self.supported_modalities}"

)

handler = self._modality_handlers[modality]

return handler(inputs)

def get_embedding_dim(self, modality: str) -> Optional[int]:

"""

Return the embedding dimension for a given modality.

For "text", this queries the SentenceTransformer model's dimension

(which triggers lazy model loading).

Args:

modality: The modality to query (e.g. "text", "image").

Returns:

The embedding dimension, or None if unknown.

"""

if modality == "text":

return self.text_model.get_sentence_embedding_dimension()

elif modality == "image":

return self.image_model.config.vision_config.hidden_size

return None

# Text Embedding

@property

def text_model(self):

if self._text_model is None:

from sentence_transformers import SentenceTransformer

self._text_model = SentenceTransformer(self.text_model_name)

return self._text_model

def _embed_text(self, inputs: List[str]) -> "np.ndarray":

return self.text_model.encode(

inputs,

batch_size=self.config.batch_size,

show_progress_bar=self.config.show_progress,

)

# Image Embedding

@property

def image_model(self):

if self._image_model is None:

from transformers import CLIPModel

self._image_model = CLIPModel.from_pretrained(self.image_model_name)

return self._image_model

@property

def image_processor(self):

if self._image_processor is None:

from transformers import CLIPProcessor

self._image_processor = CLIPProcessor.from_pretrained(self.image_model_name)

return self._image_processor

def _embed_image(self, inputs: List[Any]) -> "np.ndarray":

from pathlib import Path

import numpy as np

from PIL import Image

all_embeddings: List["np.ndarray"] = []

batch_size = self.config.batch_size

for start in range(0, len(inputs), batch_size):

batch = inputs[start : start + batch_size]

images = []

opened: List[Image.Image] = []

try:

for inp in batch:

if isinstance(

inp, (str, Path)

): # If the input string path is too large that It gives error and we could not open the image.

img = Image.open(inp)

opened.append(img)

images.append(img)

else:

images.append(inp)

processed = self.image_processor(images=images, return_tensors="pt")

finally:

for opened_img in opened:

opened_img.close()

embeddings = self.image_model.get_image_features(**processed)

embeddings = embeddings / embeddings.norm(p=2, dim=-1, keepdim=True)

all_embeddings.append(embeddings.detach().numpy())

return np.concatenate(all_embeddings, axis=0)