orca-sdk 0.0.96__py3-none-any.whl → 0.0.98__py3-none-any.whl

orca_sdk/embedding_model.py

@@ -2,28 +2,20 @@ from __future__ import annotations
 
 from abc import abstractmethod
 from datetime import datetime
-from typing import TYPE_CHECKING, Literal, Sequence, cast, overload
-
-from ._generated_api_client.api import (
-    create_finetuned_embedding_model,
-    delete_finetuned_embedding_model,
-    embed_with_finetuned_model_gpu,
-    embed_with_pretrained_model_gpu,
-    get_finetuned_embedding_model,
-    get_pretrained_embedding_model,
-    list_finetuned_embedding_models,
-    list_pretrained_embedding_models,
-)
-from ._generated_api_client.models import (
+from typing import TYPE_CHECKING, Literal, Sequence, cast, get_args, overload
+
+from ._shared.metrics import ClassificationMetrics, RegressionMetrics
+from ._utils.common import UNSET, CreateMode, DropMode
+from .client import (
+    EmbeddingEvaluationRequest,
     EmbeddingFinetuningMethod,
     EmbedRequest,
     FinetunedEmbeddingModelMetadata,
     FinetuneEmbeddingModelRequest,
-    FinetuneEmbeddingModelRequestTrainingArgs,
     PretrainedEmbeddingModelMetadata,
     PretrainedEmbeddingModelName,
+    orca_api,
 )
-from ._utils.common import CreateMode, DropMode
 from .datasource import Datasource
 from .job import Job, Status
 
@@ -32,52 +24,218 @@ if TYPE_CHECKING:
 
 
 class _EmbeddingModel:
-    name: str
     embedding_dim: int
     max_seq_length: int
     uses_context: bool
+    supports_instructions: bool
 
-    def __init__(self, *, name: str, embedding_dim: int, max_seq_length: int, uses_context: bool):
-        self.name = name
+    def __init__(
+        self, *, name: str, embedding_dim: int, max_seq_length: int, uses_context: bool, supports_instructions: bool
+    ):
         self.embedding_dim = embedding_dim
         self.max_seq_length = max_seq_length
         self.uses_context = uses_context
+        self.supports_instructions = supports_instructions
 
     @classmethod
     @abstractmethod
     def all(cls) -> Sequence[_EmbeddingModel]:
         pass
 
+    def _get_instruction_error_message(self) -> str:
+        """Get error message for instruction not supported"""
+        if isinstance(self, FinetunedEmbeddingModel):
+            return f"Model {self.name} does not support instructions. Instruction-following is only supported by models based on instruction-supporting models."
+        elif isinstance(self, PretrainedEmbeddingModel):
+            return f"Model {self.name} does not support instructions. Instruction-following is only supported by instruction-supporting models."
+        else:
+            raise ValueError("Invalid embedding model")
+
     @overload
-    def embed(self, value: str, max_seq_length: int | None = None) -> list[float]:
+    def embed(self, value: str, max_seq_length: int | None = None, prompt: str | None = None) -> list[float]:
         pass
 
     @overload
-    def embed(self, value: list[str], max_seq_length: int | None = None) -> list[list[float]]:
+    def embed(
+        self, value: list[str], max_seq_length: int | None = None, prompt: str | None = None
+    ) -> list[list[float]]:
         pass
 
-    def embed(self, value: str | list[str], max_seq_length: int | None = None) -> list[float] | list[list[float]]:
+    def embed(
+        self, value: str | list[str], max_seq_length: int | None = None, prompt: str | None = None
+    ) -> list[float] | list[list[float]]:
         """
         Generate embeddings for a value or list of values
 
         Params:
             value: The value or list of values to embed
             max_seq_length: The maximum sequence length to truncate the input to
+            prompt: Optional prompt for prompt-following embedding models.
 
         Returns:
             A matrix of floats representing the embedding for each value if the input is a list of
                 values, or a list of floats representing the embedding for the single value if the
                 input is a single value
         """
-        request = EmbedRequest(values=value if isinstance(value, list) else [value], max_seq_length=max_seq_length)
+        payload: EmbedRequest = {
+            "values": value if isinstance(value, list) else [value],
+            "max_seq_length": max_seq_length,
+            "prompt": prompt,
+        }
         if isinstance(self, PretrainedEmbeddingModel):
-            embeddings = embed_with_pretrained_model_gpu(self._model_name, body=request)
+            embeddings = orca_api.POST(
+                "/gpu/pretrained_embedding_model/{model_name}/embedding",
+                params={"model_name": cast(PretrainedEmbeddingModelName, self.name)},
+                json=payload,
+                timeout=30,  # may be slow in case of cold start
+            )
         elif isinstance(self, FinetunedEmbeddingModel):
-            embeddings = embed_with_finetuned_model_gpu(self.id, body=request)
+            embeddings = orca_api.POST(
+                "/gpu/finetuned_embedding_model/{name_or_id}/embedding",
+                params={"name_or_id": self.id},
+                json=payload,
+                timeout=30,  # may be slow in case of cold start
+            )
         else:
             raise ValueError("Invalid embedding model")
         return embeddings if isinstance(value, list) else embeddings[0]
 
+    @overload
+    def evaluate(
+        self,
+        datasource: Datasource,
+        *,
+        value_column: str = "value",
+        label_column: str,
+        score_column: None = None,
+        eval_datasource: Datasource | None = None,
+        subsample: int | None = None,
+        neighbor_count: int = 5,
+        batch_size: int = 32,
+        weigh_memories: bool = True,
+        background: Literal[True],
+    ) -> Job[ClassificationMetrics]:
+        pass
+
+    @overload
+    def evaluate(
+        self,
+        datasource: Datasource,
+        *,
+        value_column: str = "value",
+        label_column: str,
+        score_column: None = None,
+        eval_datasource: Datasource | None = None,
+        subsample: int | None = None,
+        neighbor_count: int = 5,
+        batch_size: int = 32,
+        weigh_memories: bool = True,
+        background: Literal[False] = False,
+    ) -> ClassificationMetrics:
+        pass
+
+    @overload
+    def evaluate(
+        self,
+        datasource: Datasource,
+        *,
+        value_column: str = "value",
+        label_column: None = None,
+        score_column: str,
+        eval_datasource: Datasource | None = None,
+        subsample: int | None = None,
+        neighbor_count: int = 5,
+        batch_size: int = 32,
+        weigh_memories: bool = True,
+        background: Literal[True],
+    ) -> Job[RegressionMetrics]:
+        pass
+
+    @overload
+    def evaluate(
+        self,
+        datasource: Datasource,
+        *,
+        value_column: str = "value",
+        label_column: None = None,
+        score_column: str,
+        eval_datasource: Datasource | None = None,
+        subsample: int | None = None,
+        neighbor_count: int = 5,
+        batch_size: int = 32,
+        weigh_memories: bool = True,
+        background: Literal[False] = False,
+    ) -> RegressionMetrics:
+        pass
+
+    def evaluate(
+        self,
+        datasource: Datasource,
+        *,
+        value_column: str = "value",
+        label_column: str | None = None,
+        score_column: str | None = None,
+        eval_datasource: Datasource | None = None,
+        subsample: int | None = None,
+        neighbor_count: int = 5,
+        batch_size: int = 32,
+        weigh_memories: bool = True,
+        background: bool = False,
+    ) -> (
+        ClassificationMetrics
+        | RegressionMetrics
+        | Job[ClassificationMetrics]
+        | Job[RegressionMetrics]
+        | Job[ClassificationMetrics | RegressionMetrics]
+    ):
+        """
+        Evaluate the finetuned embedding model
+        """
+        payload: EmbeddingEvaluationRequest = {
+            "datasource_name_or_id": datasource.id,
+            "datasource_label_column": label_column,
+            "datasource_value_column": value_column,
+            "datasource_score_column": score_column,
+            "eval_datasource_name_or_id": eval_datasource.id if eval_datasource is not None else None,
+            "subsample": subsample,
+            "neighbor_count": neighbor_count,
+            "batch_size": batch_size,
+            "weigh_memories": weigh_memories,
+        }
+        if isinstance(self, PretrainedEmbeddingModel):
+            response = orca_api.POST(
+                "/pretrained_embedding_model/{model_name}/evaluation",
+                params={"model_name": self.name},
+                json=payload,
+            )
+        elif isinstance(self, FinetunedEmbeddingModel):
+            response = orca_api.POST(
+                "/finetuned_embedding_model/{name_or_id}/evaluation",
+                params={"name_or_id": self.id},
+                json=payload,
+            )
+        else:
+            raise ValueError("Invalid embedding model")
+
+        def get_result(task_id: str) -> ClassificationMetrics | RegressionMetrics:
+            if isinstance(self, PretrainedEmbeddingModel):
+                res = orca_api.GET(
+                    "/pretrained_embedding_model/{model_name}/evaluation/{task_id}",
+                    params={"model_name": self.name, "task_id": task_id},
+                )["result"]
+            elif isinstance(self, FinetunedEmbeddingModel):
+                res = orca_api.GET(
+                    "/finetuned_embedding_model/{name_or_id}/evaluation/{task_id}",
+                    params={"name_or_id": self.id, "task_id": task_id},
+                )["result"]
+            else:
+                raise ValueError("Invalid embedding model")
+            assert res is not None
+            return RegressionMetrics(**res) if "mse" in res else ClassificationMetrics(**res)
+
+        job = Job(response["task_id"], lambda: get_result(response["task_id"]))
+        return job if background else job.result()
+
 
 class _ModelDescriptor:
     """
@@ -126,7 +284,7 @@ class _ModelDescriptor:
         # Load the model on first access
         if self.model is None:
             try:
-                self.model = PretrainedEmbeddingModel._get(self.name)
+                self.model = PretrainedEmbeddingModel._get(cast(PretrainedEmbeddingModelName, self.name))
             except (KeyError, AttributeError):
                 raise AttributeError(f"No embedding model named {self.name}")
 
@@ -152,17 +310,27 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
     - **`GIST_LARGE`**: GIST-Large embedding model from Hugging Face ([avsolatorio/GIST-large-Embedding-v0](https://huggingface.co/avsolatorio/GIST-large-Embedding-v0))
     - **`MXBAI_LARGE`**: Mixbreas's Large embedding model from Hugging Face ([mixedbread-ai/mxbai-embed-large-v1](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1))
     - **`QWEN2_1_5B`**: Alibaba's Qwen2-1.5B instruction-tuned embedding model from Hugging Face ([Alibaba-NLP/gte-Qwen2-1.5B-instruct](https://huggingface.co/Alibaba-NLP/gte-Qwen2-1.5B-instruct))
+    - **`BGE_BASE`**: BAAI's BGE-Base instruction-tuned embedding model from Hugging Face ([BAAI/bge-base-en-v1.5](https://huggingface.co/BAAI/bge-base-en-v1.5))
 
+    **Instruction Support:**
+
+    Some models support instruction-following for better task-specific embeddings. You can check if a model supports instructions
+    using the `supports_instructions` attribute.
 
     Examples:
         >>> PretrainedEmbeddingModel.CDE_SMALL
         PretrainedEmbeddingModel({name: CDE_SMALL, embedding_dim: 768, max_seq_length: 512})
 
+        >>> # Using instruction with an instruction-supporting model
+        >>> model = PretrainedEmbeddingModel.E5_LARGE
+        >>> embeddings = model.embed("Hello world", prompt="Represent this sentence for retrieval:")
+
     Attributes:
         name: Name of the pretrained embedding model
         embedding_dim: Dimension of the embeddings that are generated by the model
         max_seq_length: Maximum input length (in tokens not characters) that this model can process. Inputs that are longer will be truncated during the embedding process
         uses_context: Whether the pretrained embedding model uses context
+        supports_instructions: Whether this model supports instruction-following
     """
 
     # Define descriptors for model access with IDE autocomplete
@@ -175,17 +343,21 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
     GIST_LARGE = _ModelDescriptor("GIST_LARGE")
     MXBAI_LARGE = _ModelDescriptor("MXBAI_LARGE")
     QWEN2_1_5B = _ModelDescriptor("QWEN2_1_5B")
+    BGE_BASE = _ModelDescriptor("BGE_BASE")
 
-    _model_name: PretrainedEmbeddingModelName
+    name: PretrainedEmbeddingModelName
 
     def __init__(self, metadata: PretrainedEmbeddingModelMetadata):
         # for internal use only, do not document
-        self._model_name = metadata.name
+        self.name = metadata["name"]
         super().__init__(
-            name=metadata.name.value,
-            embedding_dim=metadata.embedding_dim,
-            max_seq_length=metadata.max_seq_length,
-            uses_context=metadata.uses_context,
+            name=metadata["name"],
+            embedding_dim=metadata["embedding_dim"],
+            max_seq_length=metadata["max_seq_length"],
+            uses_context=metadata["uses_context"],
+            supports_instructions=(
+                bool(metadata["supports_instructions"]) if "supports_instructions" in metadata else False
+            ),
         )
 
     def __eq__(self, other) -> bool:
@@ -202,19 +374,24 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         Returns:
             A list of all pretrained embedding models available in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_pretrained_embedding_models()]
+        return [cls(metadata) for metadata in orca_api.GET("/pretrained_embedding_model")]
 
     _instances: dict[str, PretrainedEmbeddingModel] = {}
 
     @classmethod
-    def _get(cls, name: PretrainedEmbeddingModelName | str) -> PretrainedEmbeddingModel:
+    def _get(cls, name: PretrainedEmbeddingModelName) -> PretrainedEmbeddingModel:
         # for internal use only, do not document - we want people to use dot notation to get the model
-        if str(name) not in cls._instances:
-            cls._instances[str(name)] = cls(get_pretrained_embedding_model(cast(PretrainedEmbeddingModelName, name)))
-        return cls._instances[str(name)]
+        cache_key = str(name)
+        if cache_key not in cls._instances:
+            metadata = orca_api.GET(
+                "/pretrained_embedding_model/{model_name}",
+                params={"model_name": name},
+            )
+            cls._instances[cache_key] = cls(metadata)
+        return cls._instances[cache_key]
 
     @classmethod
-    def open(cls, name: str) -> PretrainedEmbeddingModel:
+    def open(cls, name: PretrainedEmbeddingModelName) -> PretrainedEmbeddingModel:
         """
         Open an embedding model by name.
 
@@ -231,9 +408,9 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
             >>> model = PretrainedEmbeddingModel.open("GTE_BASE")
         """
         try:
-            # Use getattr to access the descriptor which will initialize the model
-            return getattr(cls, name)
-        except AttributeError:
+            # Always use the _get method which handles caching properly
+            return cls._get(name)
+        except (KeyError, AttributeError):
             raise ValueError(f"Unknown model name: {name}")
 
     @classmethod
@@ -247,7 +424,7 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         Returns:
             True if the pretrained embedding model exists, False otherwise
         """
-        return name in PretrainedEmbeddingModelName
+        return name in get_args(PretrainedEmbeddingModelName)
 
     @overload
     def finetune(
@@ -258,7 +435,7 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod | str = EmbeddingFinetuningMethod.CLASSIFICATION,
+        training_method: EmbeddingFinetuningMethod = "classification",
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: Literal[True],
@@ -274,7 +451,7 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod | str = EmbeddingFinetuningMethod.CLASSIFICATION,
+        training_method: EmbeddingFinetuningMethod = "classification",
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: Literal[False] = False,
@@ -289,7 +466,7 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         eval_datasource: Datasource | None = None,
         label_column: str = "label",
         value_column: str = "value",
-        training_method: EmbeddingFinetuningMethod | str = EmbeddingFinetuningMethod.CLASSIFICATION,
+        training_method: EmbeddingFinetuningMethod = "classification",
         training_args: dict | None = None,
         if_exists: CreateMode = "error",
         background: bool = False,
@@ -329,32 +506,35 @@ class PretrainedEmbeddingModel(_EmbeddingModel):
         elif exists and if_exists == "open":
             existing = FinetunedEmbeddingModel.open(name)
 
-            if existing.base_model_name != self._model_name:
+            if existing.base_model_name != self.name:
                 raise ValueError(f"Finetuned embedding model '{name}' already exists, but with different base model")
 
             return existing
 
         from .memoryset import LabeledMemoryset
 
-        train_datasource_id = train_datasource.id if isinstance(train_datasource, Datasource) else None
-        train_memoryset_id = train_datasource.id if isinstance(train_datasource, LabeledMemoryset) else None
-        assert train_datasource_id is not None or train_memoryset_id is not None
-        res = create_finetuned_embedding_model(
-            body=FinetuneEmbeddingModelRequest(
-                name=name,
-                base_model=self._model_name,
-                train_memoryset_id=train_memoryset_id,
-                train_datasource_id=train_datasource_id,
-                eval_datasource_id=eval_datasource.id if eval_datasource is not None else None,
-                label_column=label_column,
-                value_column=value_column,
-                training_method=EmbeddingFinetuningMethod(training_method),
-                training_args=(FinetuneEmbeddingModelRequestTrainingArgs.from_dict(training_args or {})),
-            ),
+        payload: FinetuneEmbeddingModelRequest = {
+            "name": name,
+            "base_model": self.name,
+            "label_column": label_column,
+            "value_column": value_column,
+            "training_method": training_method,
+            "training_args": training_args or {},
+        }
+        if isinstance(train_datasource, Datasource):
+            payload["train_datasource_name_or_id"] = train_datasource.id
+        elif isinstance(train_datasource, LabeledMemoryset):
+            payload["train_memoryset_name_or_id"] = train_datasource.id
+        if eval_datasource is not None:
+            payload["eval_datasource_name_or_id"] = eval_datasource.id
+
+        res = orca_api.POST(
+            "/finetuned_embedding_model",
+            json=payload,
         )
         job = Job(
-            res.finetuning_task_id,
-            lambda: FinetunedEmbeddingModel.open(res.id),
+            res["finetuning_task_id"],
+            lambda: FinetunedEmbeddingModel.open(res["id"]),
         )
         return job if background else job.result()
 
@@ -374,22 +554,27 @@ class FinetunedEmbeddingModel(_EmbeddingModel):
     """
 
     id: str
+    name: str
     created_at: datetime
     updated_at: datetime
+    base_model_name: PretrainedEmbeddingModelName
     _status: Status
 
     def __init__(self, metadata: FinetunedEmbeddingModelMetadata):
         # for internal use only, do not document
-        self.id = metadata.id
-        self.created_at = metadata.created_at
-        self.updated_at = metadata.updated_at
-        self.base_model_name = metadata.base_model
-        self._status = Status(metadata.finetuning_status.value)
+        self.id = metadata["id"]
+        self.name = metadata["name"]
+        self.created_at = datetime.fromisoformat(metadata["created_at"])
+        self.updated_at = datetime.fromisoformat(metadata["updated_at"])
+        self.base_model_name = metadata["base_model"]
+        self._status = Status(metadata["finetuning_status"])
+
         super().__init__(
-            name=metadata.name,
-            embedding_dim=metadata.embedding_dim,
-            max_seq_length=metadata.max_seq_length,
-            uses_context=metadata.uses_context,
+            name=metadata["name"],
+            embedding_dim=metadata["embedding_dim"],
+            max_seq_length=metadata["max_seq_length"],
+            uses_context=metadata["uses_context"],
+            supports_instructions=self.base_model.supports_instructions,
         )
 
     def __eq__(self, other) -> bool:
@@ -401,7 +586,7 @@ class FinetunedEmbeddingModel(_EmbeddingModel):
             f"    name: {self.name},\n"
             f"    embedding_dim: {self.embedding_dim},\n"
             f"    max_seq_length: {self.max_seq_length},\n"
-            f"    base_model: PretrainedEmbeddingModel.{self.base_model_name.value}\n"
+            f"    base_model: PretrainedEmbeddingModel.{self.base_model_name}\n"
             "})"
         )
 
@@ -418,7 +603,7 @@ class FinetunedEmbeddingModel(_EmbeddingModel):
         Returns:
             A list of all finetuned embedding model handles in the OrcaCloud
         """
-        return [cls(metadata) for metadata in list_finetuned_embedding_models()]
+        return [cls(metadata) for metadata in orca_api.GET("/finetuned_embedding_model")]
 
     @classmethod
     def open(cls, name: str) -> FinetunedEmbeddingModel:
@@ -434,7 +619,11 @@ class FinetunedEmbeddingModel(_EmbeddingModel):
         Raises:
             LookupError: If the finetuned embedding model does not exist
         """
-        return cls(get_finetuned_embedding_model(name))
+        metadata = orca_api.GET(
+            "/finetuned_embedding_model/{name_or_id}",
+            params={"name_or_id": name},
+        )
+        return cls(metadata)
 
     @classmethod
     def exists(cls, name_or_id: str) -> bool:
@@ -465,7 +654,10 @@ class FinetunedEmbeddingModel(_EmbeddingModel):
             LookupError: If the finetuned embedding model does not exist and `if_not_exists` is `"error"`
         """
         try:
-            delete_finetuned_embedding_model(name_or_id)
+            orca_api.DELETE(
+                "/finetuned_embedding_model/{name_or_id}",
+                params={"name_or_id": name_or_id},
+            )
         except (LookupError, RuntimeError):
             if if_not_exists == "error":
                 raise

orca_sdk/embedding_model_test.py

@@ -1,10 +1,12 @@
 import logging
+from typing import get_args
 from uuid import uuid4
 
 import pytest
 
 from .datasource import Datasource
 from .embedding_model import (
+    ClassificationMetrics,
     FinetunedEmbeddingModel,
     PretrainedEmbeddingModel,
     PretrainedEmbeddingModelName,
@@ -30,16 +32,16 @@ def test_open_pretrained_model_unauthenticated(unauthenticated):
 
 def test_open_pretrained_model_not_found():
     with pytest.raises(LookupError):
-        PretrainedEmbeddingModel._get("INVALID_MODEL")
+        PretrainedEmbeddingModel._get("INVALID_MODEL")  # type: ignore
 
 
 def test_all_pretrained_models():
     models = PretrainedEmbeddingModel.all()
     assert len(models) > 1
-    if len(models) != len(PretrainedEmbeddingModelName):
+    if len(models) != len(get_args(PretrainedEmbeddingModelName)):
         logging.warning("Please regenerate the SDK client! Some pretrained model names are not exposed yet.")
     model_names = [m.name for m in models]
-    assert all(enum_member in model_names for enum_member in PretrainedEmbeddingModelName.__members__)
+    assert all(m in model_names for m in get_args(PretrainedEmbeddingModelName))
 
 
 def test_embed_text():
@@ -55,6 +57,13 @@ def test_embed_text_unauthenticated(unauthenticated):
         PretrainedEmbeddingModel.GTE_BASE.embed("I love this airline", max_seq_length=32)
 
 
+def test_evaluate_pretrained_model(datasource: Datasource):
+    metrics = PretrainedEmbeddingModel.GTE_BASE.evaluate(datasource=datasource, label_column="label")
+    assert metrics is not None
+    assert isinstance(metrics, ClassificationMetrics)
+    assert metrics.accuracy > 0.5
+
+
 @pytest.fixture(scope="session")
 def finetuned_model(datasource) -> FinetunedEmbeddingModel:
     return PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource)
@@ -83,18 +92,14 @@ def test_finetune_model_with_memoryset(readonly_memoryset: LabeledMemoryset):
 
 def test_finetune_model_already_exists_error(datasource: Datasource, finetuned_model):
     with pytest.raises(ValueError):
-        PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource, value_column="text")
+        PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource)
 
 
 def test_finetune_model_already_exists_return(datasource: Datasource, finetuned_model):
     with pytest.raises(ValueError):
-        PretrainedEmbeddingModel.GTE_BASE.finetune(
-            "test_finetuned_model", datasource, if_exists="open", value_column="text"
-        )
+        PretrainedEmbeddingModel.GTE_BASE.finetune("test_finetuned_model", datasource, if_exists="open")
 
-    new_model = PretrainedEmbeddingModel.DISTILBERT.finetune(
-        "test_finetuned_model", datasource, if_exists="open", value_column="text"
-    )
+    new_model = PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model", datasource, if_exists="open")
     assert new_model is not None
     assert new_model.name == "test_finetuned_model"
     assert new_model.base_model == PretrainedEmbeddingModel.DISTILBERT
@@ -105,9 +110,7 @@ def test_finetune_model_already_exists_return(datasource: Datasource, finetuned_
 
 def test_finetune_model_unauthenticated(unauthenticated, datasource: Datasource):
     with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.DISTILBERT.finetune(
-            "test_finetuned_model_unauthenticated", datasource, value_column="text"
-        )
+        PretrainedEmbeddingModel.DISTILBERT.finetune("test_finetuned_model_unauthenticated", datasource)
 
 
 def test_use_finetuned_model_in_memoryset(datasource: Datasource, finetuned_model: FinetunedEmbeddingModel):
@@ -166,7 +169,7 @@ def test_drop_finetuned_model(datasource: Datasource):
 
 def test_drop_finetuned_model_unauthenticated(unauthenticated, datasource: Datasource):
     with pytest.raises(ValueError, match="Invalid API key"):
-        PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource, value_column="text")
+        PretrainedEmbeddingModel.DISTILBERT.finetune("finetuned_model_to_delete", datasource)
 
 
 def test_drop_finetuned_model_not_found():
@@ -179,3 +182,18 @@ def test_drop_finetuned_model_not_found():
 def test_drop_finetuned_model_unauthorized(unauthorized, finetuned_model: FinetunedEmbeddingModel):
     with pytest.raises(LookupError):
         FinetunedEmbeddingModel.drop(finetuned_model.id)
+
+
+def test_supports_instructions():
+    model = PretrainedEmbeddingModel.GTE_BASE
+    assert not model.supports_instructions
+
+    instruction_model = PretrainedEmbeddingModel.BGE_BASE
+    assert instruction_model.supports_instructions
+
+
+def test_use_explicit_instruction_prompt():
+    model = PretrainedEmbeddingModel.BGE_BASE
+    assert model.supports_instructions
+    input = "Hello world"
+    assert model.embed(input, prompt="Represent this sentence for sentiment retrieval:") != model.embed(input)