feat: specify text embedding dim

Author: aaron-ang

daft/ai/lm_studio/protocols/text_embedder.py

@@ -6,7 +6,7 @@
 from openai import OpenAI
 
 from daft import DataType
-from daft.ai.openai.protocols.text_embedder import OpenAITextEmbedder, get_input_text_token_limit_for_model
+from daft.ai.openai.protocols.text_embedder import OpenAITextEmbedder, _models, get_input_text_token_limit_for_model
 from daft.ai.protocols import TextEmbedder, TextEmbedderDescriptor
 from daft.ai.typing import EmbeddingDimensions, EmbedTextOptions, Options, UDFOptions
 from daft.utils import from_dict
@@ -26,10 +26,19 @@ class LMStudioTextEmbedderDescriptor(TextEmbedderDescriptor):
     provider_name: str
     provider_options: OpenAIProviderOptions
     model_name: str
+    dimensions: int | None = None
     embed_options: EmbedTextOptions = field(
         default_factory=lambda: EmbedTextOptions(batch_size=64, max_retries=3, on_error="raise")
     )
 
+    def __post_init__(self) -> None:
+        if self.dimensions is None:
+            return
+        if self.model_name in _models and not _models[self.model_name].supports_overriding_dimensions:
+            raise ValueError(f"Embedding model '{self.model_name}' does not support specifying dimensions")
+        if "supports_overriding_dimensions" not in self.embed_options:
+            self.embed_options["supports_overriding_dimensions"] = True
+
     def get_provider(self) -> str:
         return "lm_studio"
 
@@ -48,6 +57,8 @@ def is_async(self) -> bool:
         return True
 
     def get_dimensions(self) -> EmbeddingDimensions:
+        if self.dimensions is not None:
+            return EmbeddingDimensions(size=self.dimensions, dtype=DataType.float32())
         try:
             client = OpenAI(**self.provider_options)
             response = client.embeddings.create(
@@ -72,6 +83,7 @@ def instantiate(self) -> TextEmbedder:
             provider_options=self.provider_options,
             model=self.model_name,
             embed_options=self.embed_options,
+            dimensions=self.dimensions if self.embed_options.get("supports_overriding_dimensions", False) else None,
             provider_name=self.get_provider(),
             batch_token_limit=batch_token_limit,
             input_text_token_limit=input_text_token_limit,

daft/ai/lm_studio/provider.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import sys
-import warnings
 from typing import TYPE_CHECKING
 
 if sys.version_info < (3, 11):
@@ -50,14 +49,10 @@ def get_text_embedder(
             LMStudioTextEmbedderDescriptor,
         )
 
-        if dimensions is not None:
-            warnings.warn(
-                f"embed_text dimensions was specified but provider {self.name} currently ignores this property: see https://github.com/Eventual-Inc/Daft/issues/5555"
-            )
-
         return LMStudioTextEmbedderDescriptor(
             provider_name=self._name,
             provider_options=self._options,
             model_name=(model or self.DEFAULT_TEXT_EMBEDDER),
+            dimensions=dimensions,
             embed_options=options,
         )

daft/ai/transformers/protocols/text_embedder.py

@@ -25,8 +25,20 @@
 @dataclass
 class TransformersTextEmbedderDescriptor(TextEmbedderDescriptor):
     model: str
+    dimensions: int | None = None
     embed_options: EmbedTextOptions = field(default_factory=lambda: EmbedTextOptions(batch_size=64))
 
+    def __post_init__(self) -> None:
+        if self.dimensions is None:
+            return
+        if self.dimensions <= 0:
+            raise ValueError("Embedding dimensions must be a positive integer.")
+        dimensions = AutoConfig.from_pretrained(self.model, trust_remote_code=True).hidden_size
+        if self.dimensions > dimensions:
+            raise ValueError(
+                f"Requested dimensions ({self.dimensions}) exceeds model output size ({dimensions}) for '{self.model}'."
+            )
+
     def get_provider(self) -> str:
         return "transformers"
 
@@ -37,6 +49,8 @@ def get_options(self) -> Options:
         return dict(self.embed_options)
 
     def get_dimensions(self) -> EmbeddingDimensions:
+        if self.dimensions is not None:
+            return EmbeddingDimensions(size=self.dimensions, dtype=DataType.float32())
         dimensions = AutoConfig.from_pretrained(self.model, trust_remote_code=True).hidden_size
         return EmbeddingDimensions(size=dimensions, dtype=DataType.float32())
 
@@ -48,20 +62,26 @@ def get_udf_options(self) -> UDFOptions:
         return udf_options
 
     def instantiate(self) -> TextEmbedder:
-        return TransformersTextEmbedder(self.model, **self.embed_options)
+        return TransformersTextEmbedder(self.model, dimensions=self.dimensions, **self.embed_options)
 
 
 class TransformersTextEmbedder(TextEmbedder):
     model: SentenceTransformer
     embed_options: EmbedTextOptions
 
-    def __init__(self, model_name_or_path: str, **embed_options: Unpack[EmbedTextOptions]):
+    def __init__(
+        self,
+        model_name_or_path: str,
+        dimensions: int | None = None,
+        **embed_options: Unpack[EmbedTextOptions],
+    ):
         # Let SentenceTransformer handle device selection automatically.
         self.model = SentenceTransformer(model_name_or_path, trust_remote_code=True, backend="torch")
         self.model.eval()
         self.embed_options = embed_options
+        self.dimensions = dimensions
 
     def embed_text(self, text: list[str]) -> list[Embedding]:
         with torch.inference_mode():
-            batch = self.model.encode(text, convert_to_numpy=True)
+            batch = self.model.encode(text, convert_to_numpy=True, truncate_dim=self.dimensions)
             return list(batch)

daft/ai/transformers/provider.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import sys
-import warnings
 from typing import TYPE_CHECKING, Any
 
 if sys.version_info < (3, 11):
@@ -91,13 +90,12 @@ def get_text_embedder(
             TransformersTextEmbedderDescriptor,
         )
 
-        if dimensions is not None:
-            warnings.warn(
-                f"embed_text dimensions was specified but provider {self.name} currently ignores this property: see https://github.com/Eventual-Inc/Daft/issues/5555"
-            )
-
         embed_options: EmbedTextOptions = options
-        return TransformersTextEmbedderDescriptor(model or self.DEFAULT_TEXT_EMBEDDER, embed_options=embed_options)
+        return TransformersTextEmbedderDescriptor(
+            model=model or self.DEFAULT_TEXT_EMBEDDER,
+            dimensions=dimensions,
+            embed_options=embed_options,
+        )
 
     def get_image_classifier(
         self,