truefoundry 0.5.1rc2__tar.gz → 0.5.1rc4__tar.gz

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.5.1rc2
+Version: 0.5.1rc4
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.5.1rc2"
+version = "0.5.1rc4"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/common/auth_service_client.py

@@ -51,7 +51,7 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
         self._api_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(
@@ -128,7 +128,7 @@ class AuthServerServiceClient(AuthServiceClient):
         self._auth_server_url = url
 
     def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
-        host_arg_str = f"--host {host}" if host else "--host HOST"
+        host_arg_str = host if host else "HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
             raise Exception(

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/__init__.py

@@ -1,4 +1,9 @@
-from truefoundry.ml.autogen.client.models import LibraryName, ModelVersionEnvironment
+from truefoundry.ml.autogen.client.models import (
+    LibraryName,
+    ModelVersionEnvironment,
+    SklearnModelSchema,
+    XGBoostModelSchema,
+)
 from truefoundry.ml.enums import (
     DataSlice,
     FileFormat,
@@ -14,10 +19,7 @@ from truefoundry.ml.log_types.artifacts.artifact import (
     BlobStorageDirectory,
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory, DataDirectoryPath
-from truefoundry.ml.log_types.artifacts.model import (
-    ModelVersion,
-    infer_signature,
-)
+from truefoundry.ml.log_types.artifacts.model import ModelVersion
 from truefoundry.ml.logger import init_logger
 from truefoundry.ml.mlfoundry_api import get_client
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
@@ -37,6 +39,7 @@ from truefoundry.ml.model_framework import (
     TensorFlowFramework,
     TransformersFramework,
     XGBoostFramework,
+    sklearn_infer_schema,
 )
 
 __all__ = [
@@ -66,14 +69,16 @@ __all__ = [
     "Plot",
     "PyTorchFramework",
     "SklearnFramework",
+    "SklearnModelSchema",
     "SpaCyFramework",
     "StatsModelsFramework",
     "TensorFlowFramework",
     "TransformersFramework",
     "ViewType",
     "XGBoostFramework",
+    "XGBoostModelSchema",
     "get_client",
-    "infer_signature",
+    "sklearn_infer_schema",
 ]
 
 init_logger()

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/__init__.py

@@ -372,6 +372,7 @@ from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import
 from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagRequestDto
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -427,3 +428,6 @@ from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
 )
 from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFramework
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/models/__init__.py

@@ -343,6 +343,7 @@ from truefoundry.ml.autogen.client.models.set_experiment_tag_request_dto import
 from truefoundry.ml.autogen.client.models.set_tag_request_dto import SetTagRequestDto
 from truefoundry.ml.autogen.client.models.signed_url_dto import SignedURLDto
 from truefoundry.ml.autogen.client.models.sklearn_framework import SklearnFramework
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.ml.autogen.client.models.source import Source
 from truefoundry.ml.autogen.client.models.source1 import Source1
 from truefoundry.ml.autogen.client.models.spa_cy_framework import SpaCyFramework
@@ -398,3 +399,6 @@ from truefoundry.ml.autogen.client.models.validation_error_loc_inner import (
     ValidationErrorLocInner,
 )
 from truefoundry.ml.autogen.client.models.xg_boost_framework import XGBoostFramework
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/models/infer_method_name.py

@@ -19,7 +19,7 @@ from aenum import Enum
 
 class InferMethodName(str, Enum):
     """
-    +label=Inference method name +usage=Name of the method used for inference
+    +label=Inference Method Name +usage=Name of the method used for inference
     """
 
     """

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/models/model_version_manifest.py

@@ -53,10 +53,6 @@ class ModelVersionManifest(BaseModel):
     step: Optional[conint(strict=True, ge=0)] = Field(
         default=0, description="+label=Step"
     )
-    model_schema: Optional[Dict[str, Any]] = Field(
-        default=None,
-        description="+label=Model Schema +usage=Schema of the model +uiType=JsonInput",
-    )
     __properties = [
         "description",
         "metadata",
@@ -65,7 +61,6 @@ class ModelVersionManifest(BaseModel):
         "framework",
         "environment",
         "step",
-        "model_schema",
     ]
 
     @validator("type")
@@ -137,7 +132,6 @@ class ModelVersionManifest(BaseModel):
                 if obj.get("environment") is not None
                 else None,
                 "step": obj.get("step") if obj.get("step") is not None else 0,
-                "model_schema": obj.get("model_schema"),
             }
         )
         return _obj

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/models/sklearn_framework.py

@@ -18,10 +18,10 @@ import pprint
 import re  # noqa: F401
 from typing import Optional
 
-from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
 from truefoundry.ml.autogen.client.models.serialization_format import (
     SerializationFormat,
 )
+from truefoundry.ml.autogen.client.models.sklearn_model_schema import SklearnModelSchema
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -34,24 +34,16 @@ class SklearnFramework(BaseModel):
         default=...,
         description="+label=Type +usage=Type of the framework +value=sklearn",
     )
-    serialization_format: Optional[SerializationFormat] = Field(
-        default=None,
-        description="+label=Serialization format +usage=Serialization format used for the model",
-    )
     model_filepath: Optional[StrictStr] = Field(
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",
     )
-    infer_method_name: Optional[InferMethodName] = Field(
+    serialization_format: Optional[SerializationFormat] = Field(
         default=None,
-        description="+label=Inference method name +usage=Name of the method used for inference",
+        description="+label=Serialization format +usage=Serialization format used for the model",
     )
-    __properties = [
-        "type",
-        "serialization_format",
-        "model_filepath",
-        "infer_method_name",
-    ]
+    model_schema: Optional[SklearnModelSchema] = None
+    __properties = ["type", "model_filepath", "serialization_format", "model_schema"]
 
     @validator("type")
     def type_validate_enum(cls, value):
@@ -82,6 +74,9 @@ class SklearnFramework(BaseModel):
     def to_dict(self):
         """Returns the dictionary representation of the model using alias"""
         _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        # override the default output from truefoundry.pydantic_v1 by calling `to_dict()` of model_schema
+        if self.model_schema:
+            _dict["model_schema"] = self.model_schema.to_dict()
         return _dict
 
     @classmethod
@@ -96,9 +91,11 @@ class SklearnFramework(BaseModel):
         _obj = SklearnFramework.parse_obj(
             {
                 "type": obj.get("type"),
-                "serialization_format": obj.get("serialization_format"),
                 "model_filepath": obj.get("model_filepath"),
-                "infer_method_name": obj.get("infer_method_name"),
+                "serialization_format": obj.get("serialization_format"),
+                "model_schema": SklearnModelSchema.from_dict(obj.get("model_schema"))
+                if obj.get("model_schema") is not None
+                else None,
             }
         )
         return _obj

truefoundry-0.5.1rc4/truefoundry/ml/autogen/client/models/sklearn_model_schema.py

@@ -0,0 +1,82 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Any, Dict
+
+from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
+from truefoundry.pydantic_v1 import BaseModel, Field, conlist
+
+
+class SklearnModelSchema(BaseModel):
+    """
+    SklearnModelSchema
+    """
+
+    infer_method_name: InferMethodName = Field(
+        default=...,
+        description="+label=Inference Method Name +usage=Name of the method used for inference",
+    )
+    inputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Input Schema +usage=Schema of the input"
+    )
+    outputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Output Schema +usage=Schema of the output"
+    )
+    __properties = ["infer_method_name", "inputs", "outputs"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> SklearnModelSchema:
+        """Create an instance of SklearnModelSchema from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> SklearnModelSchema:
+        """Create an instance of SklearnModelSchema from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return SklearnModelSchema.parse_obj(obj)
+
+        _obj = SklearnModelSchema.parse_obj(
+            {
+                "infer_method_name": obj.get("infer_method_name"),
+                "inputs": obj.get("inputs"),
+                "outputs": obj.get("outputs"),
+            }
+        )
+        return _obj

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client/models/xg_boost_framework.py

@@ -21,6 +21,9 @@ from typing import Optional
 from truefoundry.ml.autogen.client.models.serialization_format import (
     SerializationFormat,
 )
+from truefoundry.ml.autogen.client.models.xg_boost_model_schema import (
+    XGBoostModelSchema,
+)
 from truefoundry.pydantic_v1 import BaseModel, Field, StrictStr, validator
 
 
@@ -41,7 +44,8 @@ class XGBoostFramework(BaseModel):
         default=None,
         description="+label=Model file path +usage=Relative path to the model file",
     )
-    __properties = ["type", "serialization_format", "model_filepath"]
+    model_schema: Optional[XGBoostModelSchema] = None
+    __properties = ["type", "serialization_format", "model_filepath", "model_schema"]
 
     @validator("type")
     def type_validate_enum(cls, value):
@@ -72,6 +76,9 @@ class XGBoostFramework(BaseModel):
     def to_dict(self):
         """Returns the dictionary representation of the model using alias"""
         _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        # override the default output from truefoundry.pydantic_v1 by calling `to_dict()` of model_schema
+        if self.model_schema:
+            _dict["model_schema"] = self.model_schema.to_dict()
         return _dict
 
     @classmethod
@@ -88,6 +95,9 @@ class XGBoostFramework(BaseModel):
                 "type": obj.get("type"),
                 "serialization_format": obj.get("serialization_format"),
                 "model_filepath": obj.get("model_filepath"),
+                "model_schema": XGBoostModelSchema.from_dict(obj.get("model_schema"))
+                if obj.get("model_schema") is not None
+                else None,
             }
         )
         return _obj

truefoundry-0.5.1rc4/truefoundry/ml/autogen/client/models/xg_boost_model_schema.py

@@ -0,0 +1,82 @@
+# coding: utf-8
+
+"""
+FastAPI
+
+No description provided (generated by Openapi Generator https://github.com/openapitools/openapi-generator)
+
+The version of the OpenAPI document: 0.1.0
+Generated by OpenAPI Generator (https://openapi-generator.tech)
+
+Do not edit the class manually.
+"""  # noqa: E501
+
+from __future__ import annotations
+
+import json
+import pprint
+import re  # noqa: F401
+from typing import Any, Dict
+
+from truefoundry.ml.autogen.client.models.infer_method_name import InferMethodName
+from truefoundry.pydantic_v1 import BaseModel, Field, conlist
+
+
+class XGBoostModelSchema(BaseModel):
+    """
+    XGBoostModelSchema
+    """
+
+    infer_method_name: InferMethodName = Field(
+        default=...,
+        description="+label=Inference Method Name +usage=Name of the method used for inference",
+    )
+    inputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Input Schema +usage=Schema of the input"
+    )
+    outputs: conlist(Dict[str, Any]) = Field(
+        default=..., description="+label= Output Schema +usage=Schema of the output"
+    )
+    __properties = ["infer_method_name", "inputs", "outputs"]
+
+    class Config:
+        """Pydantic configuration"""
+
+        allow_population_by_field_name = True
+        validate_assignment = True
+
+    def to_str(self) -> str:
+        """Returns the string representation of the model using alias"""
+        return pprint.pformat(self.dict(by_alias=True))
+
+    def to_json(self) -> str:
+        """Returns the JSON representation of the model using alias"""
+        return json.dumps(self.to_dict())
+
+    @classmethod
+    def from_json(cls, json_str: str) -> XGBoostModelSchema:
+        """Create an instance of XGBoostModelSchema from a JSON string"""
+        return cls.from_dict(json.loads(json_str))
+
+    def to_dict(self):
+        """Returns the dictionary representation of the model using alias"""
+        _dict = self.dict(by_alias=True, exclude={}, exclude_none=True)
+        return _dict
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> XGBoostModelSchema:
+        """Create an instance of XGBoostModelSchema from a dict"""
+        if obj is None:
+            return None
+
+        if not isinstance(obj, dict):
+            return XGBoostModelSchema.parse_obj(obj)
+
+        _obj = XGBoostModelSchema.parse_obj(
+            {
+                "infer_method_name": obj.get("infer_method_name"),
+                "inputs": obj.get("inputs"),
+                "outputs": obj.get("outputs"),
+            }
+        )
+        return _obj

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/client_README.md

@@ -304,6 +304,7 @@ Class | Method | HTTP request | Description
  - [SetTagRequestDto](truefoundry/ml/autogen/client/docs/SetTagRequestDto.md)
  - [SignedURLDto](truefoundry/ml/autogen/client/docs/SignedURLDto.md)
  - [SklearnFramework](truefoundry/ml/autogen/client/docs/SklearnFramework.md)
+ - [SklearnModelSchema](truefoundry/ml/autogen/client/docs/SklearnModelSchema.md)
  - [Source](truefoundry/ml/autogen/client/docs/Source.md)
  - [Source1](truefoundry/ml/autogen/client/docs/Source1.md)
  - [SpaCyFramework](truefoundry/ml/autogen/client/docs/SpaCyFramework.md)
@@ -331,6 +332,7 @@ Class | Method | HTTP request | Description
  - [ValidationError](truefoundry/ml/autogen/client/docs/ValidationError.md)
  - [ValidationErrorLocInner](truefoundry/ml/autogen/client/docs/ValidationErrorLocInner.md)
  - [XGBoostFramework](truefoundry/ml/autogen/client/docs/XGBoostFramework.md)
+ - [XGBoostModelSchema](truefoundry/ml/autogen/client/docs/XGBoostModelSchema.md)
 
 
 <a id="documentation-for-authorization"></a>

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/entities/artifacts.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  artifacts.json
-#   timestamp: 2024-12-04T12:00:28+00:00
+#   timestamp: 2024-12-05T14:45:34+00:00
 
 from __future__ import annotations
 
@@ -60,6 +60,19 @@ class BaseArtifactVersion(BaseModel):
     )
 
 
+class BaseModelSchema(BaseModel):
+    infer_method_name: str = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
+    )
+    inputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Input Schema\n+usage=Schema of the input"
+    )
+    outputs: List[Dict[str, Any]] = Field(
+        ..., description="+label= Output Schema\n+usage=Schema of the output"
+    )
+
+
 class MimeType(str, Enum):
     """
     +label=MIME Type
@@ -289,7 +302,7 @@ class SerializationFormat(str, Enum):
 
 class InferMethodName(str, Enum):
     """
-    +label=Inference method name
+    +label=Inference Method Name
     +usage=Name of the method used for inference
     """
 
@@ -297,26 +310,10 @@ class InferMethodName(str, Enum):
     predict_proba = "predict_proba"
 
 
-class SklearnFramework(BaseModel):
-    """
-    +docs=Scikit-learn framework for the model version
-    +label=Sklearn
-    """
-
-    type: Literal["sklearn"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
-    )
-    serialization_format: Optional[SerializationFormat] = Field(
-        None,
-        description="+label=Serialization format\n+usage=Serialization format used for the model",
-    )
-    model_filepath: Optional[str] = Field(
-        None,
-        description="+label=Model file path\n+usage=Relative path to the model file",
-    )
-    infer_method_name: Optional[InferMethodName] = Field(
-        None,
-        description="+label=Inference method name\n+usage=Name of the method used for inference",
+class SklearnModelSchema(BaseModelSchema):
+    infer_method_name: InferMethodName = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
 
 
@@ -451,22 +448,10 @@ class UserMessage(BaseModel):
     )
 
 
-class XGBoostFramework(BaseModel):
-    """
-    +docs=XGBoost framework for the model version
-    +label=XGBoost
-    """
-
-    type: Literal["xgboost"] = Field(
-        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
-    )
-    serialization_format: Optional[SerializationFormat] = Field(
-        None,
-        description="+label=Serialization format\n+usage=Serialization format used for the model",
-    )
-    model_filepath: Optional[str] = Field(
-        None,
-        description="+label=Model file path\n+usage=Relative path to the model file",
+class XGBoostModelSchema(BaseModelSchema):
+    infer_method_name: InferMethodName = Field(
+        ...,
+        description="+label=Inference Method Name\n+usage=Name of the method used for inference",
     )
 
 
@@ -563,6 +548,53 @@ class ChatPrompt(BasePrompt):
     )
 
 
+class SklearnFramework(BaseModel):
+    """
+    +docs=Scikit-learn framework for the model version
+    +label=Sklearn
+    """
+
+    type: Literal["sklearn"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=sklearn"
+    )
+    model_filepath: Optional[str] = Field(
+        None,
+        description="+label=Model file path\n+usage=Relative path to the model file",
+    )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_schema: Optional[SklearnModelSchema] = None
+
+
+class XGBoostFramework(BaseModel):
+    """
+    +docs=XGBoost framework for the model version
+    +label=XGBoost
+    """
+
+    type: Literal["xgboost"] = Field(
+        ..., description="+label=Type\n+usage=Type of the framework\n+value=xgboost"
+    )
+    serialization_format: Optional[SerializationFormat] = Field(
+        None,
+        description="+label=Serialization format\n+usage=Serialization format used for the model",
+    )
+    model_filepath: Optional[str] = Field(
+        None,
+        description="+label=Model file path\n+usage=Relative path to the model file",
+    )
+    model_schema: Optional[XGBoostModelSchema] = None
+
+
+class AgentApp(BaseModel):
+    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
+    tools: List[AgentOpenAPIToolWithFQN]
+    agents: List[AgentWithFQN]
+    root_agent: constr(min_length=1)
+
+
 class ModelVersion(BaseArtifactVersion):
     type: Literal["model-version"] = Field(
         ..., description='+label=Type\n+usage=Model Version\n+value="model-version"'
@@ -593,17 +625,6 @@ class ModelVersion(BaseArtifactVersion):
     )
     environment: Optional[ModelVersionEnvironment] = None
     step: conint(ge=0) = Field(0, description="+label=Step")
-    model_schema: Optional[Dict[str, Any]] = Field(
-        None,
-        description="+label=Model Schema\n+usage=Schema of the model\n+uiType=JsonInput",
-    )
-
-
-class AgentApp(BaseModel):
-    type: Literal["agent-app"] = Field(..., description="+value=agent-app")
-    tools: List[AgentOpenAPIToolWithFQN]
-    agents: List[AgentWithFQN]
-    root_agent: constr(min_length=1)
 
 
 class VersionedArtifactType(BaseModel):

{truefoundry-0.5.1rc2 → truefoundry-0.5.1rc4}/truefoundry/ml/autogen/models/signature.py

@@ -1,13 +1,16 @@
 from dataclasses import dataclass, is_dataclass
-from typing import Any, Dict, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import numpy as np
-import pandas as pd
+
+if TYPE_CHECKING:
+    import pandas as pd
+
 
 from .schema import ParamSchema, Schema, convert_dataclass_to_schema
 from .utils import infer_param_schema, infer_schema
 
-MlflowInferableDataset = Union[pd.DataFrame, np.ndarray, Dict[str, np.ndarray]]
+MlflowInferableDataset = Union["pd.DataFrame", "np.ndarray", Dict[str, "np.ndarray"]]
 
 
 class ModelSignature: