orca-sdk 0.0.90__py3-none-any.whl → 0.0.91__py3-none-any.whl

orca_sdk/_generated_api_client/models/__init__.py

@@ -3,6 +3,7 @@
 from .analyze_neighbor_labels_result import AnalyzeNeighborLabelsResult
 from .api_key_metadata import ApiKeyMetadata
 from .api_key_metadata_scope_item import ApiKeyMetadataScopeItem
+from .base_label_prediction_result import BaseLabelPredictionResult
 from .base_model import BaseModel
 from .body_create_datasource_datasource_post import BodyCreateDatasourceDatasourcePost
 from .classification_evaluation_result import ClassificationEvaluationResult
@@ -19,6 +20,7 @@ from .create_labeled_memoryset_request import CreateLabeledMemorysetRequest
 from .create_rac_model_request import CreateRACModelRequest
 from .datasource_metadata import DatasourceMetadata
 from .delete_memories_request import DeleteMemoriesRequest
+from .delete_memorysets_request import DeleteMemorysetsRequest
 from .embed_request import EmbedRequest
 from .embedding_evaluation_request import EmbeddingEvaluationRequest
 from .embedding_evaluation_response import EmbeddingEvaluationResponse
@@ -41,7 +43,6 @@ from .internal_server_error_response import InternalServerErrorResponse
 from .label_class_metrics import LabelClassMetrics
 from .label_prediction_memory_lookup import LabelPredictionMemoryLookup
 from .label_prediction_memory_lookup_metadata import LabelPredictionMemoryLookupMetadata
-from .label_prediction_result import LabelPredictionResult
 from .label_prediction_with_memories_and_feedback import LabelPredictionWithMemoriesAndFeedback
 from .labeled_memory import LabeledMemory
 from .labeled_memory_insert import LabeledMemoryInsert
@@ -56,6 +57,7 @@ from .labeled_memory_with_feedback_metrics import LabeledMemoryWithFeedbackMetri
 from .labeled_memory_with_feedback_metrics_feedback_metrics import LabeledMemoryWithFeedbackMetricsFeedbackMetrics
 from .labeled_memory_with_feedback_metrics_metadata import LabeledMemoryWithFeedbackMetricsMetadata
 from .labeled_memoryset_metadata import LabeledMemorysetMetadata
+from .labeled_memoryset_update import LabeledMemorysetUpdate
 from .list_memories_request import ListMemoriesRequest
 from .list_predictions_request import ListPredictionsRequest
 from .lookup_request import LookupRequest
@@ -93,6 +95,7 @@ from .pretrained_embedding_model_metadata import PretrainedEmbeddingModelMetadat
 from .pretrained_embedding_model_name import PretrainedEmbeddingModelName
 from .rac_head_type import RACHeadType
 from .rac_model_metadata import RACModelMetadata
+from .rac_model_update import RACModelUpdate
 from .roc_curve import ROCCurve
 from .service_unavailable_error_response import ServiceUnavailableErrorResponse
 from .task import Task
@@ -112,6 +115,7 @@ __all__ = (
     "AnalyzeNeighborLabelsResult",
     "ApiKeyMetadata",
     "ApiKeyMetadataScopeItem",
+    "BaseLabelPredictionResult",
     "BaseModel",
     "BodyCreateDatasourceDatasourcePost",
     "ClassificationEvaluationResult",
@@ -128,6 +132,7 @@ __all__ = (
     "CreateRACModelRequest",
     "DatasourceMetadata",
     "DeleteMemoriesRequest",
+    "DeleteMemorysetsRequest",
     "EmbeddingEvaluationRequest",
     "EmbeddingEvaluationResponse",
     "EmbeddingEvaluationResult",
@@ -156,6 +161,7 @@ __all__ = (
     "LabeledMemoryMetadata",
     "LabeledMemoryMetrics",
     "LabeledMemorysetMetadata",
+    "LabeledMemorysetUpdate",
     "LabeledMemoryUpdate",
     "LabeledMemoryUpdateMetadataType0",
     "LabeledMemoryWithFeedbackMetrics",
@@ -163,7 +169,6 @@ __all__ = (
     "LabeledMemoryWithFeedbackMetricsMetadata",
     "LabelPredictionMemoryLookup",
     "LabelPredictionMemoryLookupMetadata",
-    "LabelPredictionResult",
     "LabelPredictionWithMemoriesAndFeedback",
     "ListMemoriesRequest",
     "ListPredictionsRequest",
@@ -202,6 +207,7 @@ __all__ = (
     "PretrainedEmbeddingModelName",
     "RACHeadType",
     "RACModelMetadata",
+    "RACModelUpdate",
     "ROCCurve",
     "ServiceUnavailableErrorResponse",
     "Task",

orca_sdk/_generated_api_client/models/{label_prediction_result.py → base_label_prediction_result.py}

@@ -15,22 +15,22 @@ from typing import Any, Type, TypeVar, Union, cast
 from attrs import define as _attrs_define
 from attrs import field as _attrs_field
 
-T = TypeVar("T", bound="LabelPredictionResult")
+T = TypeVar("T", bound="BaseLabelPredictionResult")
 
 
 @_attrs_define
-class LabelPredictionResult:
+class BaseLabelPredictionResult:
     """Predicted label and confidence for a single input.
 
     Attributes:
-        prediction_id (str):
+        prediction_id (Union[None, str]):
         confidence (float):
         anomaly_score (Union[None, float]):
         label (int):
         label_name (Union[None, str]):
     """
 
-    prediction_id: str
+    prediction_id: Union[None, str]
     confidence: float
     anomaly_score: Union[None, float]
     label: int
@@ -38,6 +38,7 @@ class LabelPredictionResult:
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
+        prediction_id: Union[None, str]
         prediction_id = self.prediction_id
 
         confidence = self.confidence
@@ -67,7 +68,13 @@ class LabelPredictionResult:
     @classmethod
     def from_dict(cls: Type[T], src_dict: dict[str, Any]) -> T:
         d = src_dict.copy()
-        prediction_id = d.pop("prediction_id")
+
+        def _parse_prediction_id(data: object) -> Union[None, str]:
+            if data is None:
+                return data
+            return cast(Union[None, str], data)
+
+        prediction_id = _parse_prediction_id(d.pop("prediction_id"))
 
         confidence = d.pop("confidence")
 
@@ -87,7 +94,7 @@ class LabelPredictionResult:
 
         label_name = _parse_label_name(d.pop("label_name"))
 
-        label_prediction_result = cls(
+        base_label_prediction_result = cls(
             prediction_id=prediction_id,
             confidence=confidence,
             anomaly_score=anomaly_score,
@@ -95,8 +102,8 @@ class LabelPredictionResult:
             label_name=label_name,
         )
 
-        label_prediction_result.additional_properties = d
-        return label_prediction_result
+        base_label_prediction_result.additional_properties = d
+        return base_label_prediction_result
 
     @property
     def additional_keys(self) -> list[str]:

orca_sdk/_generated_api_client/models/delete_memorysets_request.py

@@ -0,0 +1,70 @@
+"""
+This file is generated by the openapi-python-client tool via the generate_api_client.py script
+
+It is a customized template from the openapi-python-client tool's default template:
+https://github.com/openapi-generators/openapi-python-client/blob/861ef5622f10fc96d240dc9becb0edf94e61446c/openapi_python_client/templates/model.py.jinja
+
+The main change is:
+- Fix typing issues
+"""
+
+# flake8: noqa: C901
+
+from typing import Any, List, Type, TypeVar, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+T = TypeVar("T", bound="DeleteMemorysetsRequest")
+
+
+@_attrs_define
+class DeleteMemorysetsRequest:
+    """
+    Attributes:
+        memoryset_ids (List[str]):
+    """
+
+    memoryset_ids: List[str]
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        memoryset_ids = self.memoryset_ids
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update(
+            {
+                "memoryset_ids": memoryset_ids,
+            }
+        )
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: Type[T], src_dict: dict[str, Any]) -> T:
+        d = src_dict.copy()
+        memoryset_ids = cast(List[str], d.pop("memoryset_ids"))
+
+        delete_memorysets_request = cls(
+            memoryset_ids=memoryset_ids,
+        )
+
+        delete_memorysets_request.additional_properties = d
+        return delete_memorysets_request
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

orca_sdk/_generated_api_client/models/labeled_memoryset_update.py

@@ -0,0 +1,113 @@
+"""
+This file is generated by the openapi-python-client tool via the generate_api_client.py script
+
+It is a customized template from the openapi-python-client tool's default template:
+https://github.com/openapi-generators/openapi-python-client/blob/861ef5622f10fc96d240dc9becb0edf94e61446c/openapi_python_client/templates/model.py.jinja
+
+The main change is:
+- Fix typing issues
+"""
+
+# flake8: noqa: C901
+
+from typing import Any, List, Type, TypeVar, Union, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="LabeledMemorysetUpdate")
+
+
+@_attrs_define
+class LabeledMemorysetUpdate:
+    """
+    Attributes:
+        label_names (Union[List[str], None, Unset]):
+        description (Union[None, Unset, str]):
+    """
+
+    label_names: Union[List[str], None, Unset] = UNSET
+    description: Union[None, Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        label_names: Union[List[str], None, Unset]
+        if isinstance(self.label_names, Unset):
+            label_names = UNSET
+        elif isinstance(self.label_names, list):
+            label_names = self.label_names
+
+        else:
+            label_names = self.label_names
+
+        description: Union[None, Unset, str]
+        if isinstance(self.description, Unset):
+            description = UNSET
+        else:
+            description = self.description
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if label_names is not UNSET:
+            field_dict["label_names"] = label_names
+        if description is not UNSET:
+            field_dict["description"] = description
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: Type[T], src_dict: dict[str, Any]) -> T:
+        d = src_dict.copy()
+
+        def _parse_label_names(data: object) -> Union[List[str], None, Unset]:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            try:
+                if not isinstance(data, list):
+                    raise TypeError()
+                label_names_type_0 = cast(List[str], data)
+
+                return label_names_type_0
+            except:  # noqa: E722
+                pass
+            return cast(Union[List[str], None, Unset], data)
+
+        label_names = _parse_label_names(d.pop("label_names", UNSET))
+
+        def _parse_description(data: object) -> Union[None, Unset, str]:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(Union[None, Unset, str], data)
+
+        description = _parse_description(d.pop("description", UNSET))
+
+        labeled_memoryset_update = cls(
+            label_names=label_names,
+            description=description,
+        )
+
+        labeled_memoryset_update.additional_properties = d
+        return labeled_memoryset_update
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

orca_sdk/_generated_api_client/models/prediction_request.py

@@ -28,12 +28,14 @@ class PredictionRequest:
         expected_labels (Union[List[int], None, Unset]):
         tags (Union[Unset, List[str]]):
         memoryset_override_id (Union[None, Unset, str]):
+        disable_telemetry (Union[Unset, bool]):  Default: False.
     """
 
     input_values: List[str]
     expected_labels: Union[List[int], None, Unset] = UNSET
     tags: Union[Unset, List[str]] = UNSET
     memoryset_override_id: Union[None, Unset, str] = UNSET
+    disable_telemetry: Union[Unset, bool] = False
     additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
 
     def to_dict(self) -> dict[str, Any]:
@@ -60,6 +62,8 @@ class PredictionRequest:
         else:
             memoryset_override_id = self.memoryset_override_id
 
+        disable_telemetry = self.disable_telemetry
+
         field_dict: dict[str, Any] = {}
         field_dict.update(self.additional_properties)
         field_dict.update(
@@ -73,6 +77,8 @@ class PredictionRequest:
             field_dict["tags"] = tags
         if memoryset_override_id is not UNSET:
             field_dict["memoryset_override_id"] = memoryset_override_id
+        if disable_telemetry is not UNSET:
+            field_dict["disable_telemetry"] = disable_telemetry
 
         return field_dict
 
@@ -150,11 +156,14 @@ class PredictionRequest:
 
         memoryset_override_id = _parse_memoryset_override_id(d.pop("memoryset_override_id", UNSET))
 
+        disable_telemetry = d.pop("disable_telemetry", UNSET)
+
         prediction_request = cls(
             input_values=input_values,
             expected_labels=expected_labels,
             tags=tags,
             memoryset_override_id=memoryset_override_id,
+            disable_telemetry=disable_telemetry,
         )
 
         prediction_request.additional_properties = d

orca_sdk/_generated_api_client/models/rac_model_update.py

@@ -0,0 +1,82 @@
+"""
+This file is generated by the openapi-python-client tool via the generate_api_client.py script
+
+It is a customized template from the openapi-python-client tool's default template:
+https://github.com/openapi-generators/openapi-python-client/blob/861ef5622f10fc96d240dc9becb0edf94e61446c/openapi_python_client/templates/model.py.jinja
+
+The main change is:
+- Fix typing issues
+"""
+
+# flake8: noqa: C901
+
+from typing import Any, Type, TypeVar, Union, cast
+
+from attrs import define as _attrs_define
+from attrs import field as _attrs_field
+
+from ..types import UNSET, Unset
+
+T = TypeVar("T", bound="RACModelUpdate")
+
+
+@_attrs_define
+class RACModelUpdate:
+    """
+    Attributes:
+        description (Union[None, Unset, str]):
+    """
+
+    description: Union[None, Unset, str] = UNSET
+    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        description: Union[None, Unset, str]
+        if isinstance(self.description, Unset):
+            description = UNSET
+        else:
+            description = self.description
+
+        field_dict: dict[str, Any] = {}
+        field_dict.update(self.additional_properties)
+        field_dict.update({})
+        if description is not UNSET:
+            field_dict["description"] = description
+
+        return field_dict
+
+    @classmethod
+    def from_dict(cls: Type[T], src_dict: dict[str, Any]) -> T:
+        d = src_dict.copy()
+
+        def _parse_description(data: object) -> Union[None, Unset, str]:
+            if data is None:
+                return data
+            if isinstance(data, Unset):
+                return data
+            return cast(Union[None, Unset, str], data)
+
+        description = _parse_description(d.pop("description", UNSET))
+
+        rac_model_update = cls(
+            description=description,
+        )
+
+        rac_model_update.additional_properties = d
+        return rac_model_update
+
+    @property
+    def additional_keys(self) -> list[str]:
+        return list(self.additional_properties.keys())
+
+    def __getitem__(self, key: str) -> Any:
+        return self.additional_properties[key]
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.additional_properties[key] = value
+
+    def __delitem__(self, key: str) -> None:
+        del self.additional_properties[key]
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.additional_properties

orca_sdk/_utils/analysis_ui.py

@@ -152,7 +152,7 @@ def display_suggested_memory_relabels(memoryset: LabeledMemoryset):
                     predicted_label_name = label_names[predicted_label]
                     predicted_label_confidence = mem.metrics.get("neighbor_predicted_label_confidence", 0)
 
-                    with gr.Row(equal_height=True, variant="panel", elem_classes="white" if i % 2 == 0 else None):
+                    with gr.Row(equal_height=True, variant="panel"):
                         with gr.Column(scale=9):
                             assert isinstance(mem.value, str)
                             gr.Markdown(mem.value, label="Value", height=50)

orca_sdk/_utils/analysis_ui_style.css

@@ -1,6 +1,3 @@
-.white {
-  background-color: white;
-}
 .centered input {
   margin: auto;
 }

orca_sdk/classification_model.py

@@ -16,6 +16,7 @@ from ._generated_api_client.api import (
     list_predictions,
     predict_gpu,
     record_prediction_feedback,
+    update_model,
 )
 from ._generated_api_client.models import (
     CreateRACModelRequest,
@@ -31,9 +32,10 @@ from ._generated_api_client.models import (
 from ._generated_api_client.models import (
     RACHeadType,
     RACModelMetadata,
+    RACModelUpdate,
 )
 from ._generated_api_client.models.prediction_request import PredictionRequest
-from ._utils.common import CreateMode, DropMode
+from ._utils.common import UNSET, CreateMode, DropMode
 from ._utils.task import wait_for_task
 from .datasource import Datasource
 from .memoryset import LabeledMemoryset
@@ -270,18 +272,53 @@ class ClassificationModel:
             if if_not_exists == "error":
                 raise
 
+    def refresh(self):
+        """Refresh the model data from the OrcaCloud"""
+        self.__dict__.update(ClassificationModel.open(self.name).__dict__)
+
+    def update_metadata(self, *, description: str | None = UNSET) -> None:
+        """
+        Update editable classification model metadata properties.
+
+        Params:
+            description: Value to set for the description, defaults to `[UNSET]` if not provided.
+
+        Examples:
+            Update the description:
+            >>> model.update(description="New description")
+
+            Remove description:
+            >>> model.update(description=None)
+        """
+        update_model(self.id, body=RACModelUpdate(description=description))
+        self.refresh()
+
     @overload
     def predict(
-        self, value: list[str], expected_labels: list[int] | None = None, tags: set[str] = set()
+        self,
+        value: list[str],
+        expected_labels: list[int] | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
     ) -> list[LabelPrediction]:
         pass
 
     @overload
-    def predict(self, value: str, expected_labels: int | None = None, tags: set[str] = set()) -> LabelPrediction:
+    def predict(
+        self,
+        value: str,
+        expected_labels: int | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
+    ) -> LabelPrediction:
         pass
 
     def predict(
-        self, value: list[str] | str, expected_labels: list[int] | int | None = None, tags: set[str] = set()
+        self,
+        value: list[str] | str,
+        expected_labels: list[int] | int | None = None,
+        tags: set[str] = set(),
+        disable_telemetry: bool = False,
     ) -> list[LabelPrediction] | LabelPrediction:
         """
         Predict label(s) for the given input value(s) grounded in similar memories
@@ -290,6 +327,7 @@ class ClassificationModel:
             value: Value(s) to get predict the labels of
             expected_labels: Expected label(s) for the given input to record for model evaluation
             tags: Tags to add to the prediction(s)
+            disable_telemetry: Whether to disable telemetry for the prediction(s)
 
         Returns:
             Label prediction or list of label predictions
@@ -318,8 +356,13 @@ class ClassificationModel:
                     else [expected_labels] if expected_labels is not None else None
                 ),
                 tags=list(tags),
+                disable_telemetry=disable_telemetry,
             ),
         )
+
+        if not disable_telemetry and any(p.prediction_id is None for p in response):
+            raise RuntimeError("Failed to save prediction to database.")
+
         predictions = [
             LabelPrediction(
                 prediction_id=prediction.prediction_id,

orca_sdk/classification_model_test.py

@@ -95,6 +95,17 @@ def test_list_models_unauthorized(unauthorized, model: ClassificationModel):
     assert ClassificationModel.all() == []
 
 
+def test_update_model(model: ClassificationModel):
+    model.update_metadata(description="New description")
+    assert model.description == "New description"
+
+
+def test_update_model_no_description(model: ClassificationModel):
+    assert model.description is not None
+    model.update_metadata(description=None)
+    assert model.description is None
+
+
 def test_delete_model(memoryset: LabeledMemoryset):
     ClassificationModel.create("model_to_delete", LabeledMemoryset.open(memoryset.name))
     assert ClassificationModel.open("model_to_delete")
@@ -168,6 +179,21 @@ def test_evaluate_with_telemetry(model):
 def test_predict(model: ClassificationModel, label_names: list[str]):
     predictions = model.predict(["Do you love soup?", "Are cats cute?"])
     assert len(predictions) == 2
+    assert predictions[0].prediction_id is not None
+    assert predictions[1].prediction_id is not None
+    assert predictions[0].label == 0
+    assert predictions[0].label_name == label_names[0]
+    assert 0 <= predictions[0].confidence <= 1
+    assert predictions[1].label == 1
+    assert predictions[1].label_name == label_names[1]
+    assert 0 <= predictions[1].confidence <= 1
+
+
+def test_predict_disable_telemetry(model: ClassificationModel, label_names: list[str]):
+    predictions = model.predict(["Do you love soup?", "Are cats cute?"], disable_telemetry=True)
+    assert len(predictions) == 2
+    assert predictions[0].prediction_id is None
+    assert predictions[1].prediction_id is None
     assert predictions[0].label == 0
     assert predictions[0].label_name == label_names[0]
     assert 0 <= predictions[0].confidence <= 1

orca_sdk/conftest.py

@@ -75,6 +75,16 @@ SAMPLE_DATA = [
     {"text": "i love cats", "label": 1, "key": "val4", "score": 0.4, "source_id": "s4"},
     {"text": "everyone loves cats", "label": 1, "key": "val5", "score": 0.5, "source_id": "s5"},
     {"text": "soup is great for the winter", "label": 0, "key": "val6", "score": 0.6, "source_id": "s6"},
+    {"text": "hot soup on a rainy day!", "label": 0, "key": "val7", "score": 0.7, "source_id": "s7"},
+    {"text": "cats sleep all day", "label": 1, "key": "val8", "score": 0.8, "source_id": "s8"},
+    {"text": "homemade soup recipes", "label": 0, "key": "val9", "score": 0.9, "source_id": "s9"},
+    {"text": "cats purr when happy", "label": 1, "key": "val10", "score": 1.0, "source_id": "s10"},
+    {"text": "chicken noodle soup is classic", "label": 0, "key": "val11", "score": 1.1, "source_id": "s11"},
+    {"text": "kittens are baby cats", "label": 1, "key": "val12", "score": 1.2, "source_id": "s12"},
+    {"text": "soup can be served cold too", "label": 0, "key": "val13", "score": 1.3, "source_id": "s13"},
+    {"text": "cats have nine lives", "label": 1, "key": "val14", "score": 1.4, "source_id": "s14"},
+    {"text": "tomato soup with grilled cheese", "label": 0, "key": "val15", "score": 1.5, "source_id": "s15"},
+    {"text": "cats are independent animals", "label": 1, "key": "val16", "score": 1.6, "source_id": "s16"},
 ]
 
 
@@ -113,4 +123,6 @@ def memoryset(datasource) -> LabeledMemoryset:
 
 @pytest.fixture(scope="session")
 def model(memoryset) -> ClassificationModel:
-    return ClassificationModel.create("test_model", memoryset, num_classes=2, memory_lookup_count=3)
+    return ClassificationModel.create(
+        "test_model", memoryset, num_classes=2, memory_lookup_count=3, description="test_description"
+    )

orca_sdk/embedding_model.py

@@ -100,6 +100,8 @@ class PretrainedEmbeddingModel(_EmbeddingModel, metaclass=_PretrainedEmbeddingMo
     - **`CDE_SMALL`**: Context-aware CDE small model from Hugging Face ([jxm/cde-small-v1](https://huggingface.co/jxm/cde-small-v1))
     - **`CLIP_BASE`**: Multi-modal CLIP model from Hugging Face ([sentence-transformers/clip-ViT-L-14](https://huggingface.co/sentence-transformers/clip-ViT-L-14))
     - **`GTE_BASE`**: Alibaba's GTE model from Hugging Face ([Alibaba-NLP/gte-base-en-v1.5](https://huggingface.co/Alibaba-NLP/gte-base-en-v1.5))
+    - **`DISTILBERT`**: DistilBERT embedding model from Hugging Face ([distilbert-base-uncased](https://huggingface.co/distilbert-base-uncased))
+    - **`GTE_SMALL`**: GTE-Small embedding model from Hugging Face ([Supabase/gte-small](https://huggingface.co/Supabase/gte-small))
 
     Examples:
         >>> PretrainedEmbeddingModel.CDE_SMALL