truefoundry 0.5.1rc2__py3-none-any.whl → 0.5.1rc4__py3-none-any.whl

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/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/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/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/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/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/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/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/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/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/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/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/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:

truefoundry/ml/autogen/models/utils.py

@@ -4,7 +4,12 @@ from collections import defaultdict
 from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
-import pandas as pd
+
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+
 
 from .exceptions import MlflowException
 from .schema import (
@@ -330,7 +335,7 @@ def infer_schema(data: Any) -> Schema:  # noqa: C901
                 ]
             )
     # pandas.Series
-    elif isinstance(data, pd.Series):
+    elif pd and isinstance(data, pd.Series):
         name = getattr(data, "name", None)
         schema = Schema(
             [
@@ -342,7 +347,7 @@ def infer_schema(data: Any) -> Schema:  # noqa: C901
             ]
         )
     # pandas.DataFrame
-    elif isinstance(data, pd.DataFrame):
+    elif pd and isinstance(data, pd.DataFrame):
         schema = Schema(
             [
                 ColSpec(
@@ -473,13 +478,13 @@ def _is_none_or_nan(x):
 
 
 def _infer_required(col) -> bool:
-    if isinstance(col, (list, pd.Series)):
+    if pd and isinstance(col, (list, pd.Series)):
         return not any(_is_none_or_nan(x) for x in col)
     return not _is_none_or_nan(col)
 
 
-def _infer_pandas_column(col: pd.Series) -> DataType:
-    if not isinstance(col, pd.Series):
+def _infer_pandas_column(col: "pd.Series") -> DataType:
+    if pd and not isinstance(col, pd.Series):
         raise TypeError(f"Expected pandas.Series, got '{type(col)}'.")
     if len(col.values.shape) > 1:
         raise MlflowException(f"Expected 1d array, got array with shape {col.shape}")
@@ -496,7 +501,7 @@ def _infer_pandas_column(col: pd.Series) -> DataType:
             # For backwards compatibility, we fall back to string
             # if the provided array is of string type
             # This is for diviner test where df field is ('key2', 'key1', 'key0')
-            if pd.api.types.is_string_dtype(col):
+            if pd and pd.api.types.is_string_dtype(col):
                 return DataType.string
             raise MlflowException(
                 f"Failed to infer schema for pandas.Series {col}. Error: {e}"

truefoundry/ml/log_types/artifacts/model.py

@@ -31,6 +31,7 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     TrueFoundryArtifactSource,
     UpdateModelVersionRequestDto,
 )
+from truefoundry.ml.autogen.models import infer_signature as _infer_signature
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.artifact import BlobStorageDirectory
@@ -105,7 +106,6 @@ class ModelVersion:
         self._deleted = False
         self._description: str = ""
         self._metadata: Dict[str, Any] = {}
-        self._model_schema: Optional[Dict[str, Any]] = None
         self._environment: Optional[ModelVersionEnvironment] = None
         self._framework: Optional[ModelFrameworkType] = None
         self._set_mutable_attrs()
@@ -149,9 +149,6 @@ class ModelVersion:
         if self._model_version.manifest:
             self._description = self._model_version.manifest.description or ""
             self._metadata = copy.deepcopy(self._model_version.manifest.metadata)
-            self._model_schema = copy.deepcopy(
-                self._model_version.manifest.model_schema
-            )
             self._environment = copy.deepcopy(self._model_version.manifest.environment)
             self._framework = (
                 copy.deepcopy(self._model_version.manifest.framework.actual_instance)
@@ -161,7 +158,6 @@ class ModelVersion:
         else:
             self._description = self._model_version.description or ""
             self._metadata = copy.deepcopy(self._model_version.artifact_metadata)
-            self._model_schema = None
             self._environment = None
             self._framework = _ModelFramework.to_model_framework_type(
                 self._model_version.model_framework
@@ -235,23 +231,6 @@ class ModelVersion:
         _validate_artifact_metadata(value)
         self._metadata = copy.deepcopy(value)
 
-    @property
-    def model_schema(self) -> Optional[Dict[str, Any]]:
-        """Get model_schema for the current model"""
-        return self._model_schema
-
-    @model_schema.setter
-    def model_schema(self, value: Optional[Dict[str, Any]]):
-        """set the model_schema for current model"""
-        if not self._model_version.manifest:
-            warnings.warn(
-                message="This model version was created using an older serialization format. model_schema will not be updated",
-                category=DeprecationWarning,
-                stacklevel=2,
-            )
-            return
-        self._model_schema = copy.deepcopy(value)
-
     @property
     def environment(self) -> Optional[ModelVersionEnvironment]:
         """Get the environment details for the model"""
@@ -466,7 +445,6 @@ class ModelVersion:
         if self._model_version.manifest:
             self._model_version.manifest.description = self.description
             self._model_version.manifest.metadata = self.metadata
-            self._model_version.manifest.model_schema = self.model_schema
             self._model_version.manifest.environment = self.environment
             self._model_version.manifest.framework = (
                 Framework.from_dict(self.framework.dict()) if self.framework else None
@@ -501,7 +479,6 @@ def _log_model_version(  # noqa: C901
     progress: Optional[bool] = None,
     framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     environment: Optional[ModelVersionEnvironment] = None,
-    model_schema: Optional[Dict[str, Any]] = None,
 ) -> ModelVersion:
     if (run and mlfoundry_artifacts_api) or (not run and not mlfoundry_artifacts_api):
         raise MlFoundryException(
@@ -613,7 +590,6 @@ def _log_model_version(  # noqa: C901
         framework=Framework.from_dict(framework.dict()) if framework else None,
         environment=environment,
         step=step,
-        model_schema=model_schema,
     )
     artifact_version_response = mlfoundry_artifacts_api.finalize_artifact_version_post(
         finalize_artifact_version_request_dto=FinalizeArtifactVersionRequestDto(
@@ -636,9 +612,6 @@ def infer_signature(
     ] = None,
     params: Optional[Dict[str, Any]] = None,
 ):
-    # TODO: Importing this globally causes hard dependencies on some libraries like pandas
-    from truefoundry.ml.autogen.models import infer_signature as _infer_signature
-
     return _infer_signature(
         model_input=model_input, model_output=model_output, params=params
     )

truefoundry/ml/mlfoundry_api.py

@@ -1265,7 +1265,6 @@ class MlFoundry:
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
         environment: Optional[ModelVersionEnvironment] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1390,7 +1389,6 @@ class MlFoundry:
             progress=progress,
             framework=framework,
             environment=environment,
-            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/mlfoundry_run.py

@@ -934,7 +934,6 @@ class MlFoundryRun:
         progress: Optional[bool] = None,
         framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
         environment: Optional[ModelVersionEnvironment] = None,
-        model_schema: Optional[Dict[str, Any]] = None,
     ) -> ModelVersion:
         # TODO (chiragjn): Document mapping of framework to list of valid model save kwargs
         # TODO (chiragjn): Add more examples
@@ -1061,7 +1060,6 @@ class MlFoundryRun:
             progress=progress,
             framework=framework,
             environment=environment,
-            model_schema=model_schema,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/model_framework.py

@@ -1,8 +1,20 @@
+import json
 import os
 import warnings
 from collections import OrderedDict
 from pickle import load as pickle_load
-from typing import Any, Callable, Dict, List, Literal, Optional, Type, Union, get_args
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    Union,
+    get_args,
+)
 
 from truefoundry.common.utils import (
     get_python_version_major_minor,
@@ -10,6 +22,7 @@ from truefoundry.common.utils import (
 )
 from truefoundry.ml.autogen.client import SerializationFormat
 from truefoundry.ml.autogen.entities import artifacts as autogen_artifacts
+from truefoundry.ml.autogen.models import infer_signature
 from truefoundry.ml.enums import ModelFramework
 from truefoundry.ml.log_types.artifacts.utils import (
     get_single_file_path_if_only_one_in_directory,
@@ -17,6 +30,9 @@ from truefoundry.ml.log_types.artifacts.utils import (
 )
 from truefoundry.pydantic_v1 import BaseModel, Field
 
+if TYPE_CHECKING:
+    from sklearn.base import BaseEstimator
+
 # Map serialization format to corresponding pip packages
 SERIALIZATION_FORMAT_TO_PACKAGES_NAME_MAP = {
     SerializationFormat.JOBLIB: ["joblib"],
@@ -421,3 +437,37 @@ def auto_update_model_framework_details(
                     model_file_path=absolute_model_filepath
                 )
             )
+
+
+def sklearn_infer_schema(
+    model_input: Any,
+    model: "BaseEstimator",
+    infer_method_name: str = "predict",
+) -> autogen_artifacts.SklearnModelSchema:
+    """
+    Infer the schema of a Sklearn model.
+
+    Args:
+        model_input (Any): The input data to be used for schema inference.
+        model (Any): The Sklearn model instance.
+        infer_method_name (str): The name of the method to be used for schema inference.
+            Eg: predict (default), predict_proba
+    Returns:
+        SklearnModelSchema: The inferred schema of the Sklearn model.
+    """
+    if not hasattr(model, infer_method_name):
+        raise ValueError(
+            f"Model does not have the method '{infer_method_name}' to infer the schema."
+        )
+    model_infer_method = getattr(model, infer_method_name)
+    model_output = model_infer_method(model_input)
+
+    model_signature = infer_signature(
+        model_input=model_input, model_output=model_output
+    )
+    model_signature_json = model_signature.to_dict()
+    return autogen_artifacts.SklearnModelSchema(
+        infer_method_name=infer_method_name,
+        inputs=json.loads(model_signature_json["inputs"]),
+        outputs=json.loads(model_signature_json["outputs"]),
+    )

{truefoundry-0.5.1rc2.dist-info → truefoundry-0.5.1rc4.dist-info}/METADATA

@@ -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.dist-info → truefoundry-0.5.1rc4.dist-info}/RECORD

@@ -26,7 +26,7 @@ truefoundry/autodeploy/utils/pydantic_compat.py,sha256=hEAUy5kLjhPdzw7yGZ2iXGMXb
 truefoundry/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/cli/__main__.py,sha256=-NkhYlT3mC5MhtekueKAvCw-sWvguj0LJRpXWzvvFjc,727
 truefoundry/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/common/auth_service_client.py,sha256=tZOa0NdATnItsMeTnEnUeTZQIgUJtpU-nvLdWtB4Px8,7978
+truefoundry/common/auth_service_client.py,sha256=RRiGUqITxeVYwKZLc923zJP-61UAvFtVlMaG2HJBvXc,7940
 truefoundry/common/constants.py,sha256=OwT8CJxGDhnrfgXCiKG5d5pkGbrd1UGGY-y672Et07Y,2310
 truefoundry/common/credential_file_manager.py,sha256=1yEk1Zm2xS4G0VDFwKSZ4w0VUrcPWQ1nJnoBaz9xyKA,4251
 truefoundry/common/credential_provider.py,sha256=Aht7hFLsnyRgMR34dRbzln7dor0WYSeA8ej8ApNmnKM,4148
@@ -107,11 +107,11 @@ truefoundry/deploy/v2/lib/models.py,sha256=pSolLMTArDuYpeNsmeeS5DWliloN_iCDfZSpR
 truefoundry/deploy/v2/lib/patched_models.py,sha256=NTU8J_CwdvEuF9zNXwFyN3suOp_196Wrm75DDy5qcXE,14184
 truefoundry/deploy/v2/lib/source.py,sha256=VHKuFREiixUP40D3Mrz-TA70spu1M0RbCzl--qwlFaY,9263
 truefoundry/logger.py,sha256=u-YCNjg5HBwE70uQcpjIG64Ghos-K2ulTWaxC03BSj4,714
-truefoundry/ml/__init__.py,sha256=aPrCvFWkb_MGMliT3rBaF4CBcF0AdaymC9DFuK5J7DM,1926
+truefoundry/ml/__init__.py,sha256=VtP0B3DNFFV_1b4WHkAo_XRMvWHmyeEwDu984t4cy8s,2040
 truefoundry/ml/artifact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/artifact/truefoundry_artifact_repo.py,sha256=FksxhUpRHb9pgWZmAB16DhXqkAL6UIAPA1S3RJUApQU,46201
 truefoundry/ml/autogen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/ml/autogen/client/__init__.py,sha256=xQKSEKW6Kn-F7GRHaDKjKNjle8JTOFiMs4Liq2kh8Rc,19383
+truefoundry/ml/autogen/client/__init__.py,sha256=uf2QeBvpQi38iGSPkAmSRki2b8nJIBx-3AnJOo9_mHM,19571
 truefoundry/ml/autogen/client/api/__init__.py,sha256=NyMBxBmIzB1o5LzZZwz9LiydHB3-hPqc_sevsnY6Jrw,746
 truefoundry/ml/autogen/client/api/auth_api.py,sha256=zpWzJhUmW6HHMY_atlUf0B25k77E1kue2hmix5I5Ih0,7017
 truefoundry/ml/autogen/client/api/deprecated_api.py,sha256=mu5x_skNcwz8v1Tr6VP72-tVP7pmBV9muyKy_2NH3F0,38824
@@ -126,7 +126,7 @@ truefoundry/ml/autogen/client/api_client.py,sha256=M31IycWorn10FYS8WbO_Whaipr0tV
 truefoundry/ml/autogen/client/api_response.py,sha256=KRyvecPMXF05PaxILHZ8JHoP4rgKBjKONMgG83aU-rM,844
 truefoundry/ml/autogen/client/configuration.py,sha256=V1oaEnxt-NfpaNmp-EZpf2glovzVhM2coWYt8HBNB4M,15723
 truefoundry/ml/autogen/client/exceptions.py,sha256=XbCbDHhYT3BVejdoGNPgEa4oS56ypkwFdxk1iOc_tFY,5355
-truefoundry/ml/autogen/client/models/__init__.py,sha256=_nSy0foDUUFyyE5zSBYKz16d5BvbWzJfCCGvMraUieY,18005
+truefoundry/ml/autogen/client/models/__init__.py,sha256=iX7kGG97Ly93SQfSwFFbggcXAf2OnO9ubNPl0OhN7W4,18193
 truefoundry/ml/autogen/client/models/add_custom_metrics_to_model_version_request_dto.py,sha256=_ISDspicTGjBCYYXubKfRYYSSQVyW3AvG-jFh47-Zfc,2163
 truefoundry/ml/autogen/client/models/add_features_to_model_version_request_dto.py,sha256=OT1-98DyWNfAHz_EgD2gMX2XkrGQ4Re945fhoAl8qSE,2099
 truefoundry/ml/autogen/client/models/agent.py,sha256=lUAbs092yo9hZmqzX6atKebDNf1uEwTf5jLvjOdvoeM,3872
@@ -206,7 +206,7 @@ truefoundry/ml/autogen/client/models/h2_o_framework.py,sha256=5X-0Xbt4tpecV_q6Ww
 truefoundry/ml/autogen/client/models/http_validation_error.py,sha256=bU-uUlz_qeFba3iHdBT7cHjLu6WZ-JmImYqgZ-kCvxM,2466
 truefoundry/ml/autogen/client/models/image_content_part.py,sha256=m_LE5XqD_UD0f8jtWt0adRgz_FoLEvgZbCDTSYORnCc,2717
 truefoundry/ml/autogen/client/models/image_url.py,sha256=TN8iEdBUmtKdTFlApUodaC4o6X_Ggd3OSOKVPv7laJE,2170
-truefoundry/ml/autogen/client/models/infer_method_name.py,sha256=01flpfaOwmqXYFZnrSJCVNl8XK4Sen5slxA6zCo4G18,788
+truefoundry/ml/autogen/client/models/infer_method_name.py,sha256=Cpg3aRGrZCG_6JypR0RMvH0cUqRvpMXxr9rabUGxFuY,788
 truefoundry/ml/autogen/client/models/internal_metadata.py,sha256=6OizrIzA3QxlEnwNFXCZGOMJuUDSYHctl_30q27uBDI,6232
 truefoundry/ml/autogen/client/models/keras_framework.py,sha256=sBKDT4Dzkvr3Rf7c8YKRbtzaO_KrPNsW8tdWI_n701Q,2150
 truefoundry/ml/autogen/client/models/latest_run_log_dto.py,sha256=wE8T8bANTb9u14Jv7DNyKWekZiUAvrzvTcE_H1cRhn4,2326
@@ -247,7 +247,7 @@ truefoundry/ml/autogen/client/models/model_response_dto.py,sha256=osrTxfygkuhxWj
 truefoundry/ml/autogen/client/models/model_server.py,sha256=N5mcWmnX0KzTMtTH8A9h_FmiQcMld3B8CXD_UYeDKaA,700
 truefoundry/ml/autogen/client/models/model_version_dto.py,sha256=9Mh4d8V33gEGYr4sys5lYYDA3bHKajR3oAZLZMAy_io,7094
 truefoundry/ml/autogen/client/models/model_version_environment.py,sha256=wTayQjMen24eJCqXQugzoXLANY5NXdVsAvZSh-QIhJA,2827
-truefoundry/ml/autogen/client/models/model_version_manifest.py,sha256=1FZ9VuMp4THs9bOrVM7YWs4eDHfQTcVTQZSN5UbJClI,4906
+truefoundry/ml/autogen/client/models/model_version_manifest.py,sha256=gRsrZzfLm1cHk5cbHmYQtShXGLLxPKcxRMPDWCqWTZc,4657
 truefoundry/ml/autogen/client/models/model_version_response_dto.py,sha256=D6XOCyggxqTkbePuypqYSHYh1PYeDv7R_J32q61sDM0,2320
 truefoundry/ml/autogen/client/models/multi_part_upload_dto.py,sha256=Ckq405vud8RQmMyKCJQJBlW5iXO7Y2mlgo8eVkiMvxg,3738
 truefoundry/ml/autogen/client/models/multi_part_upload_response_dto.py,sha256=VjH0kvl7rMjgDHjYGHnsh7KsZ5-qn-k3ksdGLJ49nIM,2431
@@ -274,7 +274,8 @@ truefoundry/ml/autogen/client/models/serialization_format.py,sha256=CYq9DR1yipEG
 truefoundry/ml/autogen/client/models/set_experiment_tag_request_dto.py,sha256=nrmi_NxLD1fI2gwlpdqFSMnBS11gRkjS4_GQFHgBcXs,2118
 truefoundry/ml/autogen/client/models/set_tag_request_dto.py,sha256=IRgAdMcWBxmjNV6nZJej4pcNfLmZwrelEZ3otwt7eeE,2144
 truefoundry/ml/autogen/client/models/signed_url_dto.py,sha256=9oHoXBj07xTdc04rqOqJO3eOjQWXCyWPhfHg-6qX60w,1897
-truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=sO3mQRTGKozmkJl4BrukZpNN9bC1WeEEtuxyio4qwn0,3290
+truefoundry/ml/autogen/client/models/sklearn_framework.py,sha256=eWpv-CNyxsD9pIBKga568DpQvCn5Nn_GvNMzY1ElRW4,3424
+truefoundry/ml/autogen/client/models/sklearn_model_schema.py,sha256=RJIbkzvIDxkFTyygQLvUfeg7sCrkhwm9Euo7FaNnrLw,2526
 truefoundry/ml/autogen/client/models/source.py,sha256=ZG2-3oIs_M8mZNPjoIRPzhHqlm63vLCVQC9idUxQ0XU,5180
 truefoundry/ml/autogen/client/models/source1.py,sha256=io35b9uVzqSvOHw6RlCXxM30thN-vHm1upaSiN2i0II,5183
 truefoundry/ml/autogen/client/models/spa_cy_framework.py,sha256=Kyrvhdalnl1uQLoPMGMry2PuFGCcQ8AGp4Cx4B9AGXQ,2150
@@ -301,15 +302,16 @@ truefoundry/ml/autogen/client/models/url.py,sha256=zMyOmdVkp1ANnQnc9GrHt42xlVwES
 truefoundry/ml/autogen/client/models/user_message.py,sha256=3TEM8qH_zT3dgM197bl44uV7m20IWn6sWeLGfStsN90,2784
 truefoundry/ml/autogen/client/models/validation_error.py,sha256=mFjwoc8g2-Usu1HXZhOQKQ4TGvLy4lwCzk8dHrJ69aA,2597
 truefoundry/ml/autogen/client/models/validation_error_loc_inner.py,sha256=nThJ5Gmy8W2Wok-ZOI4sK7uRe1BAkLS0qzq-XZbq8zs,4915
-truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=_YYI6stH0_FRkB7zPfS_85hM_9biwYLr2pLhMBFjARA,2892
+truefoundry/ml/autogen/client/models/xg_boost_framework.py,sha256=ZTvdmfX5eB_ZbA7Yy62NYAzbYKOQkPtA-zF5xxRnVfE,3429
+truefoundry/ml/autogen/client/models/xg_boost_model_schema.py,sha256=5bzi5NEeR9HetiHFJwwkmeC4LE2WvGF9hp9dTLPfEJ4,2526
 truefoundry/ml/autogen/client/rest.py,sha256=9goba8qHjQuVx5O_yRaTKu7PvBnb7r7swfy3dwuTEgk,14281
-truefoundry/ml/autogen/client_README.md,sha256=9z0MWex6ROr_RjjQ_yOZj9iXWle0pMvNITY1sR_D6TM,36112
-truefoundry/ml/autogen/entities/artifacts.py,sha256=Gz7lmWcQAKNsfLr3Qb5S0-orXuzLNbX66qMEscpyEwY,19483
+truefoundry/ml/autogen/client_README.md,sha256=8M-8txeNhpwAiXx7z5UEjPYSq2SoqA1EqtCkkKZDCMc,36276
+truefoundry/ml/autogen/entities/artifacts.py,sha256=x3DxHXPS5OpNb1vmuy6p44g2hymq6BXnivFqv6mQ-ds,20113
 truefoundry/ml/autogen/models/__init__.py,sha256=--TGRea9SQBWFfwtcl3ekb1XGfFTdEkQGSG8a2SJ60I,187
 truefoundry/ml/autogen/models/exceptions.py,sha256=q3n7FGBrg_hUy1UyoefhMnhcXUAaqXlLIGHoOVzn_d8,1390
 truefoundry/ml/autogen/models/schema.py,sha256=IhpO9qbygLqEamP3NIt3m90SseJXCOm1ZTqNbNbW-M0,55772
-truefoundry/ml/autogen/models/signature.py,sha256=EppDGjws10CTatxcTBTuT-b3jJX9lmR6ZX98eC6PB5w,4942
-truefoundry/ml/autogen/models/utils.py,sha256=Z52SnAF7UL7uluNtA_VXuB9iPju-q9bpTTFJnUh0ZTQ,24289
+truefoundry/ml/autogen/models/signature.py,sha256=rBjpxUIsEeWM0sIyYG5uCJB18DKHR4k5yZw8TzuoP48,4987
+truefoundry/ml/autogen/models/utils.py,sha256=c7RtSLXhOLcP8rjuUtfnMdaKVTZvvbsmw98gPAkAFrs,24371
 truefoundry/ml/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/ml/cli/cli.py,sha256=ckBcjUpqfhgrPE1okqT_G2iouOLt-0KjpLhHp2YdVFU,256
 truefoundry/ml/cli/commands/__init__.py,sha256=diDUiRUX4l6TtNLI4iF-ZblczkELM7FRViJ-8gGNJQY,82
@@ -331,7 +333,7 @@ truefoundry/ml/log_types/artifacts/artifact.py,sha256=6X4lO23DwM7yTgUOY4RRoYMrf9
 truefoundry/ml/log_types/artifacts/constants.py,sha256=qKxQ5mMvJE4j83BvGW3qNTKunxCiBg_EEjTdgbgJtyE,1036
 truefoundry/ml/log_types/artifacts/dataset.py,sha256=a4dxd2EN8p7Ci-cLGGiDOboN3t0395_XhWE1dmTw1Q4,13112
 truefoundry/ml/log_types/artifacts/general_artifact.py,sha256=B4XErLr-m6RmQWtxMTu3wlFRFcqSwPYp6J0OL4Ng6L0,3179
-truefoundry/ml/log_types/artifacts/model.py,sha256=BqZPUt3rKYdlceZ35NNU6a31m0FVEKYbrHQb9tIe8wg,23854
+truefoundry/ml/log_types/artifacts/model.py,sha256=sOqMptIpr2YRv7VC6AntFDvDUg5eQLsS7jQ_4TUB5Ds,22724
 truefoundry/ml/log_types/artifacts/utils.py,sha256=yhIZYe3Y-J0rG0Wcwql_Qjxe9hlefFbXUfRFqg4E9cM,7502
 truefoundry/ml/log_types/image/__init__.py,sha256=fcOq8yQnNj1rkLcPeIjLXBpdA1WIeiPsXOlAAvMxx7M,76
 truefoundry/ml/log_types/image/constants.py,sha256=wLtGEOA4T5fZHSlOXPuNDLX3lpbCtwlvGKPFk_1fah0,255
@@ -342,9 +344,9 @@ truefoundry/ml/log_types/plot.py,sha256=HuYvvRA5r8V0xAIuuqMME2IHb9d3SfGHUiuEkOP3
 truefoundry/ml/log_types/pydantic_base.py,sha256=eBlw_AEyAz4iJKDP4zgJOCFWcldwQqpf7FADW1jzIQY,272
 truefoundry/ml/log_types/utils.py,sha256=xjJ21jdPScvFmw3TbVh5NCzbzJwaqiXJyiiT4xxX1EI,335
 truefoundry/ml/logger.py,sha256=VT-BF3BnBYTWVq87O58F0c8uXMu94gYzsiFlGY3_7Ao,458
-truefoundry/ml/mlfoundry_api.py,sha256=HbtMh8uFFueJmxcmYTohV3GnziqbSGkOaq3Nv_I-1eU,62458
-truefoundry/ml/mlfoundry_run.py,sha256=C4fhMc9o1m_ubkRsR8R8rk888ktpqWd0pkqZZQ-Pr4k,44691
-truefoundry/ml/model_framework.py,sha256=2JTQ8m5oqy2FaIBqIEgPeDEfwNBfcQHTw-gdcMXF1pA,15033
+truefoundry/ml/mlfoundry_api.py,sha256=PefznXwBPBv8TD4RYid9aiE8WWEkZQAHoaAbvklFoRs,62364
+truefoundry/ml/mlfoundry_run.py,sha256=RLTZGz3htyxNmG1Xat9WevawEIjDoUCbWrSNql8WynI,44597
+truefoundry/ml/model_framework.py,sha256=I04pbczUh6bBmxWI-fzWgpUSY_ab7Js2zDB6fX4zvYM,16458
 truefoundry/ml/run_utils.py,sha256=0W208wSLUrbdfk2pjNcZlkUi9bNxG2JORqoe-5rVqHI,2423
 truefoundry/ml/session.py,sha256=F83GTC5WwGBjnJ69Ct8MqMnlutYc56JCc6YhEY1Wl-A,5394
 truefoundry/ml/validation_utils.py,sha256=J5atNhcJLvKj64ralSV9Y5Fv1Rt4SE237ICdP9-7sP4,12149
@@ -365,7 +367,7 @@ truefoundry/workflow/remote_filesystem/tfy_signed_url_client.py,sha256=5mBCIc-ON
 truefoundry/workflow/remote_filesystem/tfy_signed_url_fs.py,sha256=Hf6Dk6Fu6P7DqsK5ULgraf9DStjgigf-kjaRAMBW-RU,8680
 truefoundry/workflow/task.py,sha256=ToitYiKcNzFCtOVQwz1W8sRjbR97eVS7vQBdbgUQtKg,1779
 truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
-truefoundry-0.5.1rc2.dist-info/METADATA,sha256=vL0TH3dHp17F31KFqk_0vWHF2xefUVPUJHN54CsI7MA,2887
-truefoundry-0.5.1rc2.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.5.1rc2.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.5.1rc2.dist-info/RECORD,,
+truefoundry-0.5.1rc4.dist-info/METADATA,sha256=Ue5yEWML5oNGt2BGDDCq3YHVdj5hP1smJEUv3hMLbT8,2887
+truefoundry-0.5.1rc4.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+truefoundry-0.5.1rc4.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.5.1rc4.dist-info/RECORD,,