orca-sdk 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

orca_sdk/async_client.py

@@ -62,10 +62,12 @@ class ActionRecommendation(TypedDict):
 class AddMemorySuggestion(TypedDict):
     value: str
     label_name: str
+    similarity: NotRequired[float | None]
 
 
 class AliveResponse(TypedDict):
     ok: bool
+    checks: dict[str, bool]
 
 
 class ApiKeyMetadata(TypedDict):
@@ -292,7 +294,7 @@ class JobStatusInfo(TypedDict):
 class LabelClassMetrics(TypedDict):
     label: int | None
     label_name: NotRequired[str | None]
-    average_lookup_score: float
+    average_lookup_score: float | None
     memory_count: int
 
 
@@ -346,7 +348,7 @@ class MemoryMetrics(TypedDict):
     cluster: NotRequired[int]
     embedding_2d: NotRequired[list]
     anomaly_score: NotRequired[float]
-    neighbor_label_logits: NotRequired[list[float]]
+    neighbor_label_logits: NotRequired[list[float] | None]
     neighbor_predicted_label: NotRequired[int | None]
     neighbor_predicted_label_ambiguity: NotRequired[float]
     neighbor_predicted_label_confidence: NotRequired[float]
@@ -555,16 +557,7 @@ class PredictiveModelUpdate(TypedDict):
 
 
 PretrainedEmbeddingModelName = Literal[
-    "CLIP_BASE",
-    "GTE_BASE",
-    "CDE_SMALL",
-    "DISTILBERT",
-    "GTE_SMALL",
-    "MXBAI_LARGE",
-    "E5_LARGE",
-    "QWEN2_1_5B",
-    "BGE_BASE",
-    "GIST_LARGE",
+    "CLIP_BASE", "GTE_BASE", "CDE_SMALL", "DISTILBERT", "GTE_SMALL", "MXBAI_LARGE", "E5_LARGE", "BGE_BASE", "GIST_LARGE"
 ]
 
 
@@ -1166,6 +1159,9 @@ class FieldValidationError(TypedDict):
 
 class AddMemoryRecommendations(TypedDict):
     memories: list[AddMemorySuggestion]
+    attempts_used: NotRequired[int]
+    partial: NotRequired[bool]
+    rejection_counts: NotRequired[dict[str, int]]
 
 
 class BootstrapClassificationModelRequest(TypedDict):
@@ -1175,7 +1171,14 @@ class BootstrapClassificationModelRequest(TypedDict):
     num_examples_per_label: NotRequired[int]
 
 
-class BootstrapClassificationModelResult(TypedDict):
+class BootstrapLabeledMemoryDataInput(TypedDict):
+    model_description: str
+    label_names: list[str]
+    initial_examples: NotRequired[list[LabeledExample]]
+    num_examples_per_label: NotRequired[int]
+
+
+class BootstrapLabeledMemoryDataResult(TypedDict):
     model_description: str
     label_names: list[str]
     model_name: str
@@ -1274,10 +1277,19 @@ class CreateClassificationModelRequest(TypedDict):
     num_classes: NotRequired[int | None]
 
 
-class CreateMemorysetRequest(TypedDict):
+class CreateMemorysetFromDatasourceRequest(TypedDict):
     name: str
     description: NotRequired[str | None]
     notes: NotRequired[str | None]
+    pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
+    finetuned_embedding_model_name_or_id: NotRequired[str | None]
+    max_seq_length_override: NotRequired[int | None]
+    label_names: NotRequired[list[str] | None]
+    index_type: NotRequired[Literal["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "DISKANN"]]
+    index_params: NotRequired[dict[str, int | float | str]]
+    prompt: NotRequired[str]
+    hidden: NotRequired[bool]
+    memory_type: NotRequired[MemoryType | None]
     datasource_name_or_id: str
     datasource_label_column: NotRequired[str | None]
     datasource_score_column: NotRequired[str | None]
@@ -1285,6 +1297,14 @@ class CreateMemorysetRequest(TypedDict):
     datasource_source_id_column: NotRequired[str | None]
     datasource_partition_id_column: NotRequired[str | None]
     remove_duplicates: NotRequired[bool]
+    batch_size: NotRequired[int]
+    subsample: NotRequired[int | float | None]
+
+
+class CreateMemorysetRequest(TypedDict):
+    name: str
+    description: NotRequired[str | None]
+    notes: NotRequired[str | None]
     pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
     finetuned_embedding_model_name_or_id: NotRequired[str | None]
     max_seq_length_override: NotRequired[int | None]
@@ -1293,9 +1313,7 @@ class CreateMemorysetRequest(TypedDict):
     index_params: NotRequired[dict[str, int | float | str]]
     prompt: NotRequired[str]
     hidden: NotRequired[bool]
-    batch_size: NotRequired[int]
-    subsample: NotRequired[int | float | None]
-    memory_type: NotRequired[MemoryType]
+    memory_type: NotRequired[MemoryType | None]
 
 
 class CreateRegressionModelRequest(TypedDict):
@@ -1655,8 +1673,8 @@ class MemorysetMetadata(TypedDict):
     created_at: str
     updated_at: str
     memories_updated_at: str
-    insertion_job_id: str
-    insertion_status: JobStatus
+    insertion_job_id: str | None
+    insertion_status: JobStatus | None
     metrics: MemorysetMetrics
     memory_type: MemoryType
     label_names: list[str] | None
@@ -1666,7 +1684,7 @@ class MemorysetMetadata(TypedDict):
     document_prompt_override: str | None
     query_prompt_override: str | None
     hidden: bool
-    insertion_task_id: str
+    insertion_task_id: str | None
 
 
 class PaginatedWorkerInfo(TypedDict):
@@ -1680,7 +1698,7 @@ class BootstrapClassificationModelMeta(TypedDict):
     datasource_meta: DatasourceMetadata
     memoryset_meta: MemorysetMetadata
     model_meta: ClassificationModelMetadata
-    agent_output: BootstrapClassificationModelResult
+    agent_output: BootstrapLabeledMemoryDataResult
 
 
 class BootstrapClassificationModelResponse(TypedDict):
@@ -1737,6 +1755,22 @@ class OrcaAsyncClient(AsyncClient):
     ) -> AliveResponse:
         pass
 
+    @overload
+    async def GET(
+        self,
+        path: Literal["/gpu/check/alive"],
+        *,
+        params: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> AliveResponse:
+        pass
+
     @overload
     async def GET(
         self,
@@ -1756,7 +1790,7 @@ class OrcaAsyncClient(AsyncClient):
     @overload
     async def GET(
         self,
-        path: Literal["/gpu/check/healthy"],
+        path: Literal["/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -1772,7 +1806,7 @@ class OrcaAsyncClient(AsyncClient):
     @overload
     async def GET(
         self,
-        path: Literal["/check/healthy"],
+        path: Literal["/gpu/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -2570,7 +2604,7 @@ class OrcaAsyncClient(AsyncClient):
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
     ) -> BootstrapClassificationModelResponse:
-        """Get the status of a bootstrap classification model job"""
+        """Get the status of a bootstrap labeled memory data job"""
         pass
 
     async def GET(
@@ -2897,6 +2931,26 @@ class OrcaAsyncClient(AsyncClient):
         path: Literal["/memoryset"],
         *,
         params: None = None,
+        json: CreateMemorysetFromDatasourceRequest,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> MemorysetMetadata:
+        pass
+
+    @overload
+    async def POST(
+        self,
+        path: Literal["/memoryset/empty"],
+        *,
+        params: None = None,
         json: CreateMemorysetRequest,
         data: None = None,
         files: None = None,
@@ -3292,6 +3346,32 @@ class OrcaAsyncClient(AsyncClient):
         """Get row count from a specific datasource with optional filtering."""
         pass
 
+    @overload
+    async def POST(
+        self,
+        path: Literal["/datasource/bootstrap_memory_data"],
+        *,
+        params: None = None,
+        json: BootstrapLabeledMemoryDataInput,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> BootstrapLabeledMemoryDataResult:
+        """
+        Bootstrap memory data using an AI agent.
+
+        This endpoint uses the bootstrap labeled memory data agent to generate
+        high-quality, diverse training examples for a classification model.
+        """
+        pass
+
     @overload
     async def POST(
         self,
@@ -3540,7 +3620,7 @@ class OrcaAsyncClient(AsyncClient):
         """
         Bootstrap a classification model by creating a memoryset with generated memories and a classification model.
 
-        This endpoint uses the bootstrap_classification_model agent to generate:
+        This endpoint uses the bootstrap_labeled_memory_data agent to generate:
         1. Memoryset configuration with appropriate settings
         2. Model configuration with optimal parameters
         3. High-quality training memories for each label

orca_sdk/classification_model.py

@@ -12,12 +12,10 @@ from ._utils.common import UNSET, CreateMode, DropMode
 from .async_client import OrcaAsyncClient
 from .client import (
     BootstrapClassificationModelMeta,
-    BootstrapClassificationModelResult,
-    ClassificationEvaluationRequest,
+    BootstrapLabeledMemoryDataResult,
     ClassificationModelMetadata,
     ClassificationPredictionRequest,
     OrcaClient,
-    PostClassificationModelByModelNameOrIdEvaluationParams,
     PredictiveModelUpdate,
     RACHeadType,
 )
@@ -43,7 +41,7 @@ class BootstrappedClassificationModel:
     datasource: Datasource | None
     memoryset: LabeledMemoryset | None
     classification_model: ClassificationModel | None
-    agent_output: BootstrapClassificationModelResult | None
+    agent_output: BootstrapLabeledMemoryDataResult | None
 
     def __init__(self, metadata: BootstrapClassificationModelMeta):
         self.datasource = Datasource.open(metadata["datasource_meta"]["id"])
@@ -116,13 +114,14 @@ class ClassificationModel:
         return isinstance(other, ClassificationModel) and self.id == other.id
 
     def __repr__(self):
+        memoryset_repr = self.memoryset.__repr__().replace("\n", "\n    ")
         return (
             "ClassificationModel({\n"
             f"    name: '{self.name}',\n"
             f"    head_type: {self.head_type},\n"
             f"    num_classes: {self.num_classes},\n"
             f"    memory_lookup_count: {self.memory_lookup_count},\n"
-            f"    memoryset: LabeledMemoryset.open('{self.memoryset.name}'),\n"
+            f"    memoryset: {memoryset_repr},\n"
             "})"
         )
 

orca_sdk/client.py

@@ -60,10 +60,12 @@ class ActionRecommendation(TypedDict):
 class AddMemorySuggestion(TypedDict):
     value: str
     label_name: str
+    similarity: NotRequired[float | None]
 
 
 class AliveResponse(TypedDict):
     ok: bool
+    checks: dict[str, bool]
 
 
 class ApiKeyMetadata(TypedDict):
@@ -290,7 +292,7 @@ class JobStatusInfo(TypedDict):
 class LabelClassMetrics(TypedDict):
     label: int | None
     label_name: NotRequired[str | None]
-    average_lookup_score: float
+    average_lookup_score: float | None
     memory_count: int
 
 
@@ -344,7 +346,7 @@ class MemoryMetrics(TypedDict):
     cluster: NotRequired[int]
     embedding_2d: NotRequired[list]
     anomaly_score: NotRequired[float]
-    neighbor_label_logits: NotRequired[list[float]]
+    neighbor_label_logits: NotRequired[list[float] | None]
     neighbor_predicted_label: NotRequired[int | None]
     neighbor_predicted_label_ambiguity: NotRequired[float]
     neighbor_predicted_label_confidence: NotRequired[float]
@@ -553,16 +555,7 @@ class PredictiveModelUpdate(TypedDict):
 
 
 PretrainedEmbeddingModelName = Literal[
-    "CLIP_BASE",
-    "GTE_BASE",
-    "CDE_SMALL",
-    "DISTILBERT",
-    "GTE_SMALL",
-    "MXBAI_LARGE",
-    "E5_LARGE",
-    "QWEN2_1_5B",
-    "BGE_BASE",
-    "GIST_LARGE",
+    "CLIP_BASE", "GTE_BASE", "CDE_SMALL", "DISTILBERT", "GTE_SMALL", "MXBAI_LARGE", "E5_LARGE", "BGE_BASE", "GIST_LARGE"
 ]
 
 
@@ -1164,6 +1157,9 @@ class FieldValidationError(TypedDict):
 
 class AddMemoryRecommendations(TypedDict):
     memories: list[AddMemorySuggestion]
+    attempts_used: NotRequired[int]
+    partial: NotRequired[bool]
+    rejection_counts: NotRequired[dict[str, int]]
 
 
 class BootstrapClassificationModelRequest(TypedDict):
@@ -1173,7 +1169,14 @@ class BootstrapClassificationModelRequest(TypedDict):
     num_examples_per_label: NotRequired[int]
 
 
-class BootstrapClassificationModelResult(TypedDict):
+class BootstrapLabeledMemoryDataInput(TypedDict):
+    model_description: str
+    label_names: list[str]
+    initial_examples: NotRequired[list[LabeledExample]]
+    num_examples_per_label: NotRequired[int]
+
+
+class BootstrapLabeledMemoryDataResult(TypedDict):
     model_description: str
     label_names: list[str]
     model_name: str
@@ -1272,10 +1275,19 @@ class CreateClassificationModelRequest(TypedDict):
     num_classes: NotRequired[int | None]
 
 
-class CreateMemorysetRequest(TypedDict):
+class CreateMemorysetFromDatasourceRequest(TypedDict):
     name: str
     description: NotRequired[str | None]
     notes: NotRequired[str | None]
+    pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
+    finetuned_embedding_model_name_or_id: NotRequired[str | None]
+    max_seq_length_override: NotRequired[int | None]
+    label_names: NotRequired[list[str] | None]
+    index_type: NotRequired[Literal["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "DISKANN"]]
+    index_params: NotRequired[dict[str, int | float | str]]
+    prompt: NotRequired[str]
+    hidden: NotRequired[bool]
+    memory_type: NotRequired[MemoryType | None]
     datasource_name_or_id: str
     datasource_label_column: NotRequired[str | None]
     datasource_score_column: NotRequired[str | None]
@@ -1283,6 +1295,14 @@ class CreateMemorysetRequest(TypedDict):
     datasource_source_id_column: NotRequired[str | None]
     datasource_partition_id_column: NotRequired[str | None]
     remove_duplicates: NotRequired[bool]
+    batch_size: NotRequired[int]
+    subsample: NotRequired[int | float | None]
+
+
+class CreateMemorysetRequest(TypedDict):
+    name: str
+    description: NotRequired[str | None]
+    notes: NotRequired[str | None]
     pretrained_embedding_model_name: NotRequired[PretrainedEmbeddingModelName | None]
     finetuned_embedding_model_name_or_id: NotRequired[str | None]
     max_seq_length_override: NotRequired[int | None]
@@ -1291,9 +1311,7 @@ class CreateMemorysetRequest(TypedDict):
     index_params: NotRequired[dict[str, int | float | str]]
     prompt: NotRequired[str]
     hidden: NotRequired[bool]
-    batch_size: NotRequired[int]
-    subsample: NotRequired[int | float | None]
-    memory_type: NotRequired[MemoryType]
+    memory_type: NotRequired[MemoryType | None]
 
 
 class CreateRegressionModelRequest(TypedDict):
@@ -1653,8 +1671,8 @@ class MemorysetMetadata(TypedDict):
     created_at: str
     updated_at: str
     memories_updated_at: str
-    insertion_job_id: str
-    insertion_status: JobStatus
+    insertion_job_id: str | None
+    insertion_status: JobStatus | None
     metrics: MemorysetMetrics
     memory_type: MemoryType
     label_names: list[str] | None
@@ -1664,7 +1682,7 @@ class MemorysetMetadata(TypedDict):
     document_prompt_override: str | None
     query_prompt_override: str | None
     hidden: bool
-    insertion_task_id: str
+    insertion_task_id: str | None
 
 
 class PaginatedWorkerInfo(TypedDict):
@@ -1678,7 +1696,7 @@ class BootstrapClassificationModelMeta(TypedDict):
     datasource_meta: DatasourceMetadata
     memoryset_meta: MemorysetMetadata
     model_meta: ClassificationModelMetadata
-    agent_output: BootstrapClassificationModelResult
+    agent_output: BootstrapLabeledMemoryDataResult
 
 
 class BootstrapClassificationModelResponse(TypedDict):
@@ -1735,6 +1753,22 @@ class OrcaClient(Client):
     ) -> AliveResponse:
         pass
 
+    @overload
+    def GET(
+        self,
+        path: Literal["/gpu/check/alive"],
+        *,
+        params: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> AliveResponse:
+        pass
+
     @overload
     def GET(
         self,
@@ -1754,7 +1788,7 @@ class OrcaClient(Client):
     @overload
     def GET(
         self,
-        path: Literal["/gpu/check/healthy"],
+        path: Literal["/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -1770,7 +1804,7 @@ class OrcaClient(Client):
     @overload
     def GET(
         self,
-        path: Literal["/check/healthy"],
+        path: Literal["/gpu/check/healthy"],
         *,
         params: None = None,
         parse_as: Literal["json"] = "json",
@@ -2568,7 +2602,7 @@ class OrcaClient(Client):
         timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
         extensions: RequestExtensions | None = None,
     ) -> BootstrapClassificationModelResponse:
-        """Get the status of a bootstrap classification model job"""
+        """Get the status of a bootstrap labeled memory data job"""
         pass
 
     def GET(
@@ -2895,6 +2929,26 @@ class OrcaClient(Client):
         path: Literal["/memoryset"],
         *,
         params: None = None,
+        json: CreateMemorysetFromDatasourceRequest,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> MemorysetMetadata:
+        pass
+
+    @overload
+    def POST(
+        self,
+        path: Literal["/memoryset/empty"],
+        *,
+        params: None = None,
         json: CreateMemorysetRequest,
         data: None = None,
         files: None = None,
@@ -3290,6 +3344,32 @@ class OrcaClient(Client):
         """Get row count from a specific datasource with optional filtering."""
         pass
 
+    @overload
+    def POST(
+        self,
+        path: Literal["/datasource/bootstrap_memory_data"],
+        *,
+        params: None = None,
+        json: BootstrapLabeledMemoryDataInput,
+        data: None = None,
+        files: None = None,
+        content: None = None,
+        parse_as: Literal["json"] = "json",
+        headers: HeaderTypes | None = None,
+        cookies: CookieTypes | None = None,
+        auth: AuthTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        follow_redirects: bool | UseClientDefault = USE_CLIENT_DEFAULT,
+        timeout: TimeoutTypes | UseClientDefault = USE_CLIENT_DEFAULT,
+        extensions: RequestExtensions | None = None,
+    ) -> BootstrapLabeledMemoryDataResult:
+        """
+        Bootstrap memory data using an AI agent.
+
+        This endpoint uses the bootstrap labeled memory data agent to generate
+        high-quality, diverse training examples for a classification model.
+        """
+        pass
+
     @overload
     def POST(
         self,
@@ -3538,7 +3618,7 @@ class OrcaClient(Client):
         """
         Bootstrap a classification model by creating a memoryset with generated memories and a classification model.
 
-        This endpoint uses the bootstrap_classification_model agent to generate:
+        This endpoint uses the bootstrap_labeled_memory_data agent to generate:
         1. Memoryset configuration with appropriate settings
         2. Model configuration with optimal parameters
         3. High-quality training memories for each label

orca_sdk/embedding_model.py

@@ -26,14 +26,23 @@ if TYPE_CHECKING:
 class EmbeddingModelBase(ABC):
     embedding_dim: int
     max_seq_length: int
+    num_params: int
     uses_context: bool
     supports_instructions: bool
 
     def __init__(
-        self, *, name: str, embedding_dim: int, max_seq_length: int, uses_context: bool, supports_instructions: bool
+        self,
+        *,
+        name: str,
+        embedding_dim: int,
+        max_seq_length: int,
+        num_params: int,
+        uses_context: bool,
+        supports_instructions: bool,
     ):
         self.embedding_dim = embedding_dim
         self.max_seq_length = max_seq_length
+        self.num_params = num_params
         self.uses_context = uses_context
         self.supports_instructions = supports_instructions
 
@@ -340,7 +349,6 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
     - **`E5_LARGE`**: E5-Large instruction-tuned embedding model from Hugging Face ([intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct))
     - **`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:**
@@ -360,6 +368,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         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
+        num_params: Number of parameters in the model
         uses_context: Whether the pretrained embedding model uses context
         supports_instructions: Whether this model supports instruction-following
     """
@@ -373,7 +382,6 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
     E5_LARGE = _ModelDescriptor("E5_LARGE")
     GIST_LARGE = _ModelDescriptor("GIST_LARGE")
     MXBAI_LARGE = _ModelDescriptor("MXBAI_LARGE")
-    QWEN2_1_5B = _ModelDescriptor("QWEN2_1_5B")
     BGE_BASE = _ModelDescriptor("BGE_BASE")
 
     name: PretrainedEmbeddingModelName
@@ -385,6 +393,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
             name=metadata["name"],
             embedding_dim=metadata["embedding_dim"],
             max_seq_length=metadata["max_seq_length"],
+            num_params=metadata["num_params"],
             uses_context=metadata["uses_context"],
             supports_instructions=(
                 bool(metadata["supports_instructions"]) if "supports_instructions" in metadata else False
@@ -395,7 +404,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         return isinstance(other, PretrainedEmbeddingModel) and self.name == other.name
 
     def __repr__(self) -> str:
-        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}}})"
+        return f"PretrainedEmbeddingModel({{name: {self.name}, embedding_dim: {self.embedding_dim}, max_seq_length: {self.max_seq_length}, num_params: {self.num_params/1000000:.0f}M}})"
 
     @classmethod
     def all(cls) -> list[PretrainedEmbeddingModel]:
@@ -543,7 +552,7 @@ class PretrainedEmbeddingModel(EmbeddingModelBase):
         elif exists and if_exists == "open":
             existing = FinetunedEmbeddingModel.open(name)
 
-            if existing.base_model_name != self.name:
+            if existing.base_model.name != self.name:
                 raise ValueError(f"Finetuned embedding model '{name}' already exists, but with different base model")
 
             return existing
@@ -598,7 +607,7 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
     name: str
     created_at: datetime
     updated_at: datetime
-    base_model_name: PretrainedEmbeddingModelName
+    base_model: PretrainedEmbeddingModel
     _status: Status
 
     def __init__(self, metadata: FinetunedEmbeddingModelMetadata):
@@ -607,13 +616,14 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
         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.base_model = PretrainedEmbeddingModel._get(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"],
+            num_params=self.base_model.num_params,
             uses_context=metadata["uses_context"],
             supports_instructions=self.base_model.supports_instructions,
         )
@@ -627,15 +637,10 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
             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}\n"
+            f"    base_model: PretrainedEmbeddingModel.{self.base_model.name}\n"
             "})"
         )
 
-    @property
-    def base_model(self) -> PretrainedEmbeddingModel:
-        """Pretrained model the finetuned embedding model was based on"""
-        return PretrainedEmbeddingModel._get(self.base_model_name)
-
     @classmethod
     def all(cls) -> list[FinetunedEmbeddingModel]:
         """
@@ -702,6 +707,6 @@ class FinetunedEmbeddingModel(EmbeddingModelBase):
                 "/finetuned_embedding_model/{name_or_id}",
                 params={"name_or_id": name_or_id},
             )
-        except (LookupError, RuntimeError):
+        except LookupError:
             if if_not_exists == "error":
                 raise